Fixed the code; now just need to clean up everything

thomas-schouten · Nov 19, 2024 · 5a8abaa · 5a8abaa
1 parent fb679f9
commit 5a8abaa
Show file tree

Hide file tree

Showing 21 changed files with 437 additions and 199 deletions.
diff --git a/.DS_Store b/.DS_Store
diff --git a/plato/__pycache__/globe.cpython-310.pyc b/plato/__pycache__/globe.cpython-310.pyc
diff --git a/plato/__pycache__/grids.cpython-310.pyc b/plato/__pycache__/grids.cpython-310.pyc
diff --git a/plato/__pycache__/optimisation.cpython-310.pyc b/plato/__pycache__/optimisation.cpython-310.pyc
diff --git a/plato/__pycache__/plate_torques.cpython-310.pyc b/plato/__pycache__/plate_torques.cpython-310.pyc
diff --git a/plato/__pycache__/plates.cpython-310.pyc b/plato/__pycache__/plates.cpython-310.pyc
diff --git a/plato/__pycache__/points.cpython-310.pyc b/plato/__pycache__/points.cpython-310.pyc
diff --git a/plato/__pycache__/settings.cpython-310.pyc b/plato/__pycache__/settings.cpython-310.pyc
diff --git a/plato/__pycache__/slabs.cpython-310.pyc b/plato/__pycache__/slabs.cpython-310.pyc
diff --git a/plato/__pycache__/utils_calc.cpython-310.pyc b/plato/__pycache__/utils_calc.cpython-310.pyc
diff --git a/plato/__pycache__/utils_data.cpython-310.pyc b/plato/__pycache__/utils_data.cpython-310.pyc
diff --git a/plato/globe.py b/plato/globe.py
@@ -187,6 +187,7 @@ def calculate_net_rotation(
             ages = None,
             cases = None,
             plateIDs = None,
+            PROGRESS_BAR = True,
         ):
         """
         Calculate the net rotation of the Earth's lithosphere.
@@ -198,7 +199,11 @@ def calculate_net_rotation(
         _cases = utils_data.select_cases(cases, self.settings.cases)
 
         # Calculate the net rotation of the Earth's lithosphere
-        for i, _age in enumerate(_ages):
+        for i, _age in enumerate(_tqdm(
+                _ages,
+                desc="Calculating net rotation",
+                disable=(self.settings.logger.level in [logging.INFO, logging.DEBUG] or not PROGRESS_BAR),
+            )):
             for _case in _cases:
                 # Check if plate data is provided
                 if utils_init.check_object_data(plates, Plates, _age, _case):
@@ -217,10 +222,10 @@ def calculate_net_rotation(
                     )
 
                 # Check if plate data is provided
-                if utils_init.check_object_data(plates, Plates, _age, _case):
+                if utils_init.check_object_data(points, Points, _age, _case):
                     logging.info(f"Calculating net rotation for case {_case} at age {_age} using provided Points data")
                     _points = points
-                elif utils_init.check_object_data(self.plates, Plates, _age, _case):
+                elif utils_init.check_object_data(self.points, Points, _age, _case):
                     logging.info(f"Calculating net rotation for case {_case} at age {_age} using stored Points data")
                     _points = self.points
                 else:
@@ -239,10 +244,17 @@ def calculate_net_rotation(
                 # Select plates and points data
                 selected_plates = _plates.data[_age][_case]
                 selected_points = _points.data[_age][_case]
+
+                # Filter by plateIDs
                 if plateIDs is not None:
                     selected_plates = _plates.data[_age][_case][_plates.data[_age][_case].plateID.isin(_plateIDs)]
                     selected_points = _points.data[_age][_case][_points.data[_age][_case].plateID.isin(_plateIDs)]
 
+                # Filter by minimum plate area
+                if self.settings.options[_case]["Minimum plate area"] > 0.:
+                    selected_plates = selected_plates[selected_plates.area >= self.settings.options[_case]["Minimum plate area"]]
+                    selected_points = selected_points[selected_points.plateID.isin(selected_plates.plateID)]
+
                 # Calculate net rotation
                 net_rotation_pole = utils_calc.compute_net_rotation(
                     selected_plates,
@@ -311,6 +323,7 @@ def save(
             self,
             cases: Optional[Union[str, List[str]]] = None,
             file_dir: Optional[str] = None,
+            PROGRESS_BAR: Optional[bool] = True,
         ):
         """
         Function to save 'Globe' object.
@@ -323,7 +336,11 @@ def save(
         _file_dir = self.settings.dir_path if file_dir is None else file_dir
 
         # Loop through ages
-        for _case in _tqdm(_cases, desc="Saving Globe", disable=self.settings.logger.level==logging.INFO):
+        for _case in _tqdm(
+                _cases, 
+                desc="Saving Globe", 
+                disable=(self.settings.logger.level in [logging.INFO, logging.DEBUG] or not PROGRESS_BAR)
+            ):
             utils_data.DataFrame_to_parquet(
                 self.data[_case],
                 "Globe",
@@ -337,6 +354,7 @@ def export(
             self,
             cases: Optional[Union[str, List[str]]] = None,
             file_dir: Optional[str] = None,
+            PROGRESS_BAR: Optional[bool] = True,
         ):
         """
         Function to export 'Globe' object.
@@ -349,7 +367,11 @@ def export(
         _file_dir = self.settings.dir_path if file_dir is None else file_dir
 
         # Loop through ages
-        for _case in _tqdm(_cases, desc="Saving Globe", disable=self.settings.logger.level==logging.INFO):
+        for _case in _tqdm(
+                _cases, 
+                desc="Exporting Globe", 
+                disable=(self.settings.logger.level in [logging.INFO, logging.DEBUG] or not PROGRESS_BAR)
+            ):
             utils_data.DataFrame_to_csv(
                 self.data[_case],
                 "Globe",

diff --git a/plato/grids.py b/plato/grids.py
@@ -353,73 +353,78 @@ def interpolate_data_to_grid(
         logging.info(f"{grid_type} updated:", getattr(self, grid_type))
 
     def save_all(
-        self,
-        ages: Union[None, List[int], List[float], _numpy.ndarray] = None,
-        cases: Union[None, str, List[str]] = None,
-        file_dir: Optional[str] = None,
+            self,
+            ages: Union[None, List[int], List[float], _numpy.ndarray] = None,
+            cases: Union[None, str, List[str]] = None,
+            file_dir: Optional[str] = None,
+            PROGRESS_BAR: bool = True,
         ):
         """
         Function to save all the grids
         """
         # Save seafloor grid
-        self.save_seafloor_age(ages, file_dir)
+        self.save_seafloor_age(ages, file_dir, PROGRESS_BAR)
 
         # Save sediment grid
-        self.save_sediment(ages, cases, file_dir)
+        self.save_sediment(ages, cases, file_dir, PROGRESS_BAR)
 
         # Save continental grid
-        self.save_continent(ages, cases, file_dir)
+        self.save_continent(ages, cases, file_dir, PROGRESS_BAR)
 
         # Save velocity grid
-        self.save_velocity(ages, cases, file_dir)
+        self.save_velocity(ages, cases, file_dir, PROGRESS_BAR)
 
     def save_seafloor_age(
             self,
             ages: Union[None, List[int], List[float], _numpy.ndarray] = None,
             file_dir: Optional[str] = None,
+            PROGRESS_BAR: bool = True,
         ):
         """
         Function to save the the seafloor age grid.
         """
-        self.save_grid(self.seafloor_age, "Seafloor_age", ages, None, file_dir)
+        self.save_grid(self.seafloor_age, "Seafloor_age", ages, None, file_dir, PROGRESS_BAR)
 
     def save_sediment(
             self,
             ages: Union[None, List[int], List[float], _numpy.ndarray] = None,
             cases: Union[None, str, List[str]] = None,
             file_dir: Optional[str] = None,
+            PROGRESS_BAR: bool = True,
         ):
         """
         Function to save the the sediment grid.
         """
         if self.sediment is not None:
-            self.save_grid(self.sediment, "Sediment", ages, cases, file_dir)
+            self.save_grid(self.sediment, "Sediment", ages, cases, file_dir, PROGRESS_BAR)
 
     def save_continent(
             self,
             ages: Union[None, List[int], List[float], _numpy.ndarray] = None,
             cases: Union[None, str, List[str]] = None,
             file_dir: Optional[str] = None,
+            PROGRESS_BAR: bool = True,
         ):
         """
         Function to save the the continental grid.
         """
         # Check if grids exists
         if self.continent is not None:
-            self.save_grid(self.continent, "Continent", ages, cases, file_dir)
+            self.save_grid(self.continent, "Continent", ages, cases, file_dir, PROGRESS_BAR)
 
     def save_velocity(
             self,
             ages: Union[None, List[int], List[float], _numpy.ndarray] = None,
             cases: Union[None, str, List[str]] = None,
             file_dir: Optional[str] = None,
+            PROGRESS_BAR: bool = True,
         ):
         """
         Function to save the the velocity grid.
         """
         # Check if grids exists
         if self.velocity is not None:
-            self.save_grid(self.velocity, "Velocity", ages, cases, file_dir)
+            self.save_grid(self.velocity, "Velocity", ages, cases, file_dir, PROGRESS_BAR)
 
     def save_grid(
             self,
@@ -428,6 +433,7 @@ def save_grid(
             ages: Union[None, List[int], List[float], _numpy.ndarray] = None,
             cases: Union[None, str, List[str]] = None,
             file_dir: Optional[str] = None,
+            PROGRESS_BAR: bool = True,
         ):
         """
         Function to save a grid
@@ -442,7 +448,11 @@ def save_grid(
         _file_dir = self.settings.dir_path if file_dir is None else file_dir
 
         # Loop through ages
-        for _age in _tqdm(_ages, desc=f"Saving {type} grids", disable=self.settings.logger.level==logging.INFO):
+        for _age in _tqdm(
+                _ages,
+                desc=f"Saving {type} grids",
+                disable=(self.settings.logger.level in [logging.INFO, logging.DEBUG] or not PROGRESS_BAR)
+            ):
             if data[_age] is Dict:
                 # Loop through cases
                 for _case in _cases: