.

AequilibraE · Sep 23, 2024 · 3db993a · 3db993a
1 parent d070e29
commit 3db993a
Show file tree

Hide file tree

Showing 8 changed files with 78 additions and 104 deletions.
diff --git a/.github/workflows/documentation.yml b/.github/workflows/documentation.yml
@@ -45,6 +45,11 @@ jobs:
 
     - name: Test docstrings
       run: |
+        python -m pytest --doctest-modules aequilibrae/distribution/gravity_application.py
+        python -m pytest --doctest-modules aequilibrae/distribution/gravity_calibration.py
+        python -m pytest --doctest-modules aequilibrae/distribution/ipf.py
+        python -m pytest --doctest-modules aequilibrae/matrix/aequilibrae_data.py 
+        python -m pytest --doctest-modules aequilibrae/matrix/aequilibrae_matrix.py
         python -m pytest --doctest-modules aequilibrae/log.py aequilibrae/parameters.py aequilibrae/paths/vdf.py
         python -m pytest docs/source/modeling_with_aequilibrae/project_pieces --doctest-glob=*.rst
 

diff --git a/aequilibrae/distribution/gravity_application.py b/aequilibrae/distribution/gravity_application.py
@@ -19,35 +19,32 @@ class GravityApplication:
     """Applies a synthetic gravity model.
 
     Model is an instance of SyntheticGravityModel class.
+
     Impedance is an instance of AequilibraEMatrix.
+
     Row and Column vectors are instances of AequilibraeData.
 
     .. code-block:: python
 
-        >>> import pandas as pd
-        >>> from aequilibrae import Project
-        >>> from aequilibrae.matrix import AequilibraeMatrix, AequilibraeData
+        >>> from aequilibrae.matrix import AequilibraeData
         >>> from aequilibrae.distribution import SyntheticGravityModel, GravityApplication
 
-        >>> project = Project.from_path("/tmp/test_project_ga")
+        >>> project = create_example(project_path)
 
         # We define the model we will use
         >>> model = SyntheticGravityModel()
 
         # Before adding a parameter to the model, you need to define the model functional form
-        >>> model.function = "GAMMA" # "EXPO" or "POWER"
+        # You can select one of GAMMA, EXPO or POWER.
+        >>> model.function = "GAMMA"
 
         # Only the parameter(s) applicable to the chosen functional form will have any effect
         >>> model.alpha = 0.1
         >>> model.beta = 0.0001
 
-        # Or you can load the model from a file
-        # model.load('path/to/model/file')
-
         # We load the impedance matrix
-        >>> matrix = AequilibraeMatrix()
-        >>> matrix.load('/tmp/test_project_ga/matrices/skims.omx')
-        >>> matrix.computational_view(['distance_blended'])
+        >>> matrix = project.matrices.get_matrix("skims")
+        >>> matrix.computational_view(["distance_blended"])
 
         # We create the vectors we will use
         >>> query = "SELECT zone_id, population, employment FROM zones;"
@@ -61,8 +58,11 @@ class GravityApplication:
         >>> zones = df.index.shape[0]
 
         # We create the vector database
-        >>> args = {"entries": zones, "field_names": ["productions", "attractions"],
-        ...     "data_types": [np.float64, np.float64], "memory_mode": True}
+        >>> args = {"entries": zones,
+        ...         "field_names": ["productions", "attractions"],
+        ...         "data_types": [np.float64, np.float64],
+        ...         "memory_mode": True}
+
         >>> vectors = AequilibraeData()
         >>> vectors.create_empty(**args)
 
@@ -81,20 +81,14 @@ class GravityApplication:
         ...         "model": model,
         ...         "columns": vectors,
         ...         "column_field": "attractions",
-        ...         "output": '/tmp/test_project_ga/matrices/matrix.aem',
+        ...         "output": os.path.join(project_path, 'matrices/gravity_matrix.aem'),
         ...         "nan_as_zero":True
         ...         }
         >>> gravity = GravityApplication(**args)
 
         # Solve and save the outputs
         >>> gravity.apply()
-        >>> gravity.output.export('/tmp/test_project_ga/matrices/omx_file.omx')
-
-        # To save your report into a file, you can do the following:
-        # with open('/tmp/test_project_ga/report.txt', 'w') as file:
-        #     for line in gravity.report:
-        #         file.write(f"{line}\\n")
-
+        >>> gravity.output.export(os.path.join(project_path, 'matrices/gravity_omx.omx'))
     """
 
     def __init__(self, project=None, **kwargs):

diff --git a/aequilibrae/distribution/gravity_calibration.py b/aequilibrae/distribution/gravity_calibration.py
@@ -17,42 +17,33 @@
 class GravityCalibration:
     """Calibrate a traditional gravity model
 
-    Available deterrence function forms are: 'EXPO' or 'POWER'. 'GAMMA'
+    Available deterrence function forms are: 'EXPO', 'POWER' or 'GAMMA'.
 
     .. code-block:: python
 
-        >>> from aequilibrae import Project
-        >>> from aequilibrae.matrix import AequilibraeMatrix
+        >>> from aequilibrae.utils.create_example import create_example
         >>> from aequilibrae.distribution import GravityCalibration
 
-        >>> project = Project.from_path("/tmp/test_project_gc")
+        >>> project = create_example(project_path)
 
-        # We load the impedance matrix
-        >>> matrix = AequilibraeMatrix()
-        >>> matrix.load('/tmp/test_project_gc/matrices/demand.omx')
-        >>> matrix.computational_view(['matrix'])
+        # We load the demand matrix
+        >>> demand = project.matrices.get_matrix("demand_omx")
+        >>> demand.computational_view()
 
-        # We load the impedance matrix
-        >>> impedmatrix = AequilibraeMatrix()
-        >>> impedmatrix.load('/tmp/test_project_gc/matrices/skims.omx')
-        >>> impedmatrix.computational_view(['time_final'])
+        # We load the skim matrix
+        >>> skim = project.matrices.get_matrix("skims")
+        >>> skim.computational_view(["time_final"])
 
-        # Creates the problem
-        >>> args = {"matrix": matrix,
-        ...         "impedance": impedmatrix,
+        >>> args = {"matrix": demand,
+        ...         "impedance": skim,
         ...         "row_field": "productions",
         ...         "function": 'expo',
         ...         "nan_as_zero": True}
         >>> gravity = GravityCalibration(**args)
 
         # Solve and save outputs
         >>> gravity.calibrate()
-        >>> gravity.model.save('/tmp/test_project_gc/dist_expo_model.mod')
-
-        # To save the model report in a file
-        # with open('/tmp/test_project_gc/report.txt', 'w') as f:
-        #     for line in gravity.report:
-        #         f.write(f'{line}\\n')
+        >>> gravity.model.save(os.path.join(project_path, 'dist_expo_model.mod'))
     """
 
     def __init__(self, project=None, **kwargs):

diff --git a/aequilibrae/distribution/ipf.py b/aequilibrae/distribution/ipf.py
@@ -17,7 +17,6 @@ class Ipf:
 
     .. code-block:: python
 
-        >>> from aequilibrae.utils.create_example import create_example
         >>> from aequilibrae.distribution import Ipf
         >>> from aequilibrae.matrix import AequilibraeData
 
@@ -26,9 +25,9 @@ class Ipf:
         >>> matrix = project.matrices.get_matrix("demand_omx")
         >>> matrix.computational_view()
 
-        >>> dataset_args = {"entries": matrix.zones, 
+        >>> dataset_args = {"entries": matrix.zones,
         ...                 "field_names": ["productions", "attractions"],
-        ...                 "data_types": [np.float64, np.float64], 
+        ...                 "data_types": [np.float64, np.float64],
         ...                 "memory_mode": True}
 
         >>> vectors = AequilibraeData()
@@ -40,11 +39,11 @@ class Ipf:
         # We assume that the indices would be sorted and that they would match the matrix indices
         >>> vectors.index[:] = matrix.index[:]
 
-        >>> ipf_args = {"matrix": matrix, 
-        ...             "rows": vectors, 
-        ...             "row_field": "productions", 
+        >>> ipf_args = {"matrix": matrix,
+        ...             "rows": vectors,
+        ...             "row_field": "productions",
         ...             "columns": vectors,
-        ...             "column_field": "attractions", 
+        ...             "column_field": "attractions",
         ...             "nan_as_zero": False}
 
         >>> fratar = Ipf(**ipf_args)

diff --git a/aequilibrae/matrix/aequilibrae_data.py b/aequilibrae/matrix/aequilibrae_data.py
@@ -12,7 +12,30 @@
 
 
 class AequilibraeData(object):
-    """AequilibraE dataset"""
+    """AequilibraE dataset
+    
+    .. code-block:: python
+
+        >>> from aequilibrae.matrix import AequilibraeData
+
+        >>> args = {"file_path": os.path.join(my_folder_path, "vectors.aed"),
+        ...         "entries": 3,
+        ...         "field_names": ["production"],
+        ...         "data_types": [np.float64]}
+
+        >>> vectors = AequilibraeData()
+        >>> vectors.create_empty(**args)
+
+        >>> vectors.production[:] = [1, 4, 7]
+        >>> vectors.index[:] = [1, 2, 3]
+
+        >>> vectors.export(os.path.join(my_folder_path, "vectors.csv"))
+
+        >>> reload_dataset = AequilibraeData()
+        >>> reload_dataset.load(os.path.join(my_folder_path, "vectors.aed"))
+        >>> reload_dataset.production
+        memmap([1., 4., 7.])
+    """
 
     def __init__(self):
         self.data = None
@@ -37,8 +60,8 @@ def create_empty(
         Creates a new empty dataset
 
         :Arguments:
-            **file_path** (:obj:`str`, *Optional*): Full path for the output data file. If *memory_mode* is 'false' and
-            path is missing, then the file is created in the temp folder
+            **file_path** (:obj:`str`, *Optional*): Full path for the output data file. If 'memory_mode' is ``False``
+            and path is missing, then the file is created in the temp folder
 
             **entries** (:obj:`int`, *Optional*): Number of records in the dataset. Default is 1
 
@@ -51,27 +74,6 @@ def create_empty(
 
             **memory_mode** (:obj:`bool`, *Optional*): If ``True``, dataset will be kept in memory.
             If ``False``, the dataset will be a memory-mapped numpy array
-
-        .. code-block:: python
-
-            >>> from aequilibrae.matrix import AequilibraeData, AequilibraeMatrix
-
-            >>> mat = AequilibraeMatrix()
-            >>> mat.load('/tmp/test_project/matrices/demand.omx')
-            >>> mat.computational_view()
-
-            >>> vectors = "/tmp/test_project/vectors.aed"
-
-            >>> args = {
-            ...      "file_path": vectors,
-            ...      "entries": mat.zones,
-            ...      "field_names": ["origins", "destinations"],
-            ...      "data_types": [np.float64, np.float64]
-            ... }
-
-            >>> dataset = AequilibraeData()
-            >>> dataset.create_empty(**args)
-
         """
 
         if file_path is not None or memory_mode:
@@ -145,13 +147,6 @@ def load(self, file_path):
 
         :Arguments:
             **file_path** (:obj:`str`): Full file path to the AequilibraeData to be loaded
-
-        .. code-block:: python
-
-            >>> from aequilibrae.matrix import AequilibraeData
-
-            >>> dataset = AequilibraeData()
-            >>> dataset.load("/tmp/test_project/vectors.aed")
         """
         f = open(file_path)
         self.file_path = os.path.realpath(f.name)
@@ -173,14 +168,6 @@ def export(self, file_name, table_name="aequilibrae_table"):
             **file_name** (:obj:`str`): File name with PATH and extension (csv, or sqlite3, sqlite or db)
 
             **table_name** (:obj:`str`): It only applies if you are saving to an SQLite table. Otherwise ignored
-
-        .. code-block:: python
-
-            >>> from aequilibrae.matrix import AequilibraeData
-
-            >>> dataset = AequilibraeData()
-            >>> dataset.load("/tmp/test_project/vectors.aed")
-            >>> dataset.export("/tmp/test_project/vectors.csv")
         """
 
         file_type = os.path.splitext(file_name)[1]
@@ -237,9 +224,8 @@ def random_name():
 
             >>> from aequilibrae.matrix import AequilibraeData
 
-            >>> name = AequilibraeData().random_name() # doctest: +ELLIPSIS
-
-            # This is an example of output
-            # '/tmp/Aequilibrae_data_5werr5f36-b123-asdf-4587-adfglkjhqwe.aed'
+            >>> dataset = AequilibraeData()
+            >>> dataset.random_name() # doctest: +ELLIPSIS
+            '/tmp/Aequilibrae_data_...'
         """
         return os.path.join(tempfile.gettempdir(), f"Aequilibrae_data_{uuid.uuid4()}.aed")
diff --git a/aequilibrae/matrix/aequilibrae_matrix.py b/aequilibrae/matrix/aequilibrae_matrix.py
@@ -823,7 +823,6 @@ def load(self, file_path: str):
 
         .. code-block:: python
 
-            >>> from aequilibrae.utils.create_example import create_example
             >>> from aequilibrae.matrix import AequilibraeMatrix
 
             >>> project = create_example(project_path)
@@ -943,9 +942,10 @@ def copy(
             >>> names_list = ['Car trips', 'pt trips', 'DRT trips', 'bike trips', 'walk trips']
 
             >>> mat = AequilibraeMatrix()
-            >>> mat.create_empty(file_name=os.path.join(my_folder_path, 'my_matrix.aem'), 
-            ...                  zones=zones_in_the_model, 
+            >>> mat.create_empty(file_name=os.path.join(my_folder_path, 'my_matrix.aem'),
+            ...                  zones=zones_in_the_model,
             ...                  matrix_names=names_list)
+
             >>> mat.copy(os.path.join(my_folder_path, 'copy_of_my_matrix.aem'),
             ...          cores=['bike trips', 'walk trips'],
             ...          names=['bicycle', 'walking'],
@@ -1007,7 +1007,6 @@ def rows(self) -> np.ndarray:
 
         .. code-block:: python
 
-            >>> from aequilibrae.utils.create_example import create_example
             >>> from aequilibrae.matrix import AequilibraeMatrix
 
             >>> project = create_example(project_path)
@@ -1037,7 +1036,6 @@ def columns(self) -> np.ndarray:
 
         .. code-block:: python
 
-            >>> from aequilibrae.utils.create_example import create_example
             >>> from aequilibrae.matrix import AequilibraeMatrix
 
             >>> project = create_example(project_path)
@@ -1069,8 +1067,8 @@ def nan_to_num(self):
             >>> index = np.arange(1, 4, dtype=np.int32)
 
             >>> mat = AequilibraeMatrix()
-            >>> mat.create_empty(file_name=os.path.join(my_folder_path, "matrices/nan_matrix.aem"), 
-            ...                  zones=3, 
+            >>> mat.create_empty(file_name=os.path.join(my_folder_path, "matrices/nan_matrix.aem"),
+            ...                  zones=3,
             ...                  matrix_names=["only_nan"])
             >>> mat.index[:] = index[:]
             >>> mat.matrix["only_nan"][:, :] = nan_matrix[:, :]
@@ -1125,7 +1123,7 @@ def setName(self, matrix_name: str):
             >>> from aequilibrae.matrix import AequilibraeMatrix
 
             >>> zones_in_the_model = 3317
-            
+
             >>> mat = AequilibraeMatrix()
             >>> mat.create_empty(file_name=os.path.join(my_folder_path, 'my_matrix.aem'),
             ...                  zones=zones_in_the_model,

diff --git a/conftest.py b/conftest.py
@@ -8,6 +8,8 @@
 
 import pytest
 
+import pandas as pd
+
 from aequilibrae import Project
 from aequilibrae.project.database_connection import database_connection
 from aequilibrae.transit import Transit
@@ -104,3 +106,6 @@ def transit_conn(create_gtfs_project):
 def doctest_fixtures(doctest_namespace, create_path, tmp_path_factory):
     doctest_namespace["project_path"] = str(create_path)
     doctest_namespace["my_folder_path"] = tmp_path_factory.mktemp(uuid.uuid4().hex)
+    doctest_namespace["os"] = os
+    doctest_namespace["pd"] = pd
+    doctest_namespace["create_example"] = create_example
diff --git a/docs/create_docs_data.py b/docs/create_docs_data.py
@@ -10,10 +10,6 @@
 
 project = create_example("/tmp/test_project")
 project.close()
-project = create_example("/tmp/test_project_gc")
-project.close()
-project = create_example("/tmp/test_project_ga")
-project.close()
 project = create_example("/tmp/accessing_sfalls_data")
 project.close()
 project = create_example("/tmp/accessing_nauru_data", "nauru")