Fixed most things

plato/optimisation.py

@@ -3,6 +3,7 @@
 import matplotlib.pyplot as plt
 
 # Plato libraries
+import plato.utils_data as utils_data
 from .plate_torques import PlateTorques
 
 class Optimisation():
@@ -58,6 +59,13 @@ def minimise_residual_torque(
         :return:                        The optimal slab pull coefficient, optimal viscosity, normalised residual torque, and indices of the optimal coefficients
         :rtype:                         float, float, numpy.ndarray, tuple
         """
+        # Define age and case if not provided
+        if age is None:
+            age = self.settings.ages[0]
+
+        if case is None:
+            case = self.settings.cases[0]
+
         # Set range of viscosities
         viscs = _numpy.linspace(viscosity_range[0], viscosity_range[1], grid_size)
 
@@ -73,114 +81,120 @@ def minimise_residual_torque(
         ones_grid = _numpy.ones_like(visc_grid)
 
         # Filter plates
-        selected_plates = self.plates[age][case].copy()
-        if plates_of_interest:
-            selected_plates = selected_plates[selected_plates["plateID"].isin(plates_of_interest)]
-            selected_plates = selected_plates.reset_index(drop=True)
-        else:
-            plates_of_interest = selected_plates["plateID"]
+        _data = self.plates.data[age][case].copy()
+
+        _plateIDs = utils_data.select_plateIDs(plateIDs, self.plates.data[age][case])
+
+        if plateIDs is not None:
+            _data = _data[_data["plateID"].isin(_plateIDs)]
+            # _data = _data.reset _index(drop=True)
+
+        print(f"_plateIDs length: {len(_plateIDs)}")
+        print(f"_data indices: {_data.index}")
+        # # Filter plates by minimum area
+        # if minimum_plate_area is None:
+        #     minimum_plate_area = self.settings.options[case]["Minimum plate area"]
 
-        # Filter plates by minimum area
-        if minimum_plate_area is None:
-            minimum_plate_area = self.settings.options[case]["Minimum plate area"]
-        selected_plates = selected_plates[selected_plates["area"] > minimum_plate_area]
-        selected_plates = selected_plates.reset_index(drop=True)
-        plates_of_interest = selected_plates["plateID"]
+        # _data = _data[_data["area"] > minimum_plate_area]
+        # _data = _data.reset_index(drop=True)
 
         # Get total area
-        total_area = selected_plates["area"].sum()
+        total_area = _data["area"].sum()
 
         # Initialise dictionaries and arrays to store driving and residual torques
-        if age not in self.driving_torque:
-            self.driving_torque[age] = {}
-        if age not in self.driving_torque_normalised:
-            self.driving_torque_normalised[age] = {}
-        if age not in self.residual_torque:
-            self.residual_torque[age] = {}
-        if age not in self.residual_torque_normalised:
-            self.residual_torque_normalised[age] = {}
+        # if age not in self.driving_torque:
+        #     self.driving_torque[age] = {}
+        # if age not in self.driving_torque_normalised:
+        #     self.driving_torque_normalised[age] = {}
+        # if age not in self.residual_torque:
+        #     self.residual_torque[age] = {}
+        # if age not in self.residual_torque_normalised:
+        #     self.residual_torque_normalised[age] = {}
 
-        self.driving_torque[age][case] = _numpy.zeros_like(sp_const_grid); self.driving_torque_normalised[age][case] = _numpy.zeros_like(sp_const_grid)
-        self.residual_torque[age][case] = _numpy.zeros_like(sp_const_grid); self.residual_torque_normalised[age][case] = _numpy.zeros_like(sp_const_grid)
+        driving_mag = _numpy.zeros_like(sp_const_grid); 
+        # self.driving_torque_normalised[age][case] = _numpy.zeros_like(sp_const_grid)
+        residual_mag = _numpy.zeros_like(sp_const_grid); 
+        # self.residual_torque_normalised[age][case] = _numpy.zeros_like(sp_const_grid)
 
         # Initialise dictionaries to store optimal coefficients
-        if age not in self.opt_i:
-            self.opt_i[age] = {}
-        if age not in self.opt_j:
-            self.opt_j[age] = {}
-        if age not in self.opt_sp_const:
-            self.opt_sp_const[age] = {}
-        if age not in self.opt_visc:
-            self.opt_visc[age] = {}
+        # if age not in self.opt_i:
+        #     self.opt_i[age] = {}
+        # if age not in self.opt_j:
+        #     self.opt_j[age] = {}
+        # if age not in self.opt_sp_const:
+        #     self.opt_sp_const[age] = {}
+        # if age not in self.opt_visc:
+        #     self.opt_visc[age] = {}
 
         # Get torques
-        for k, _ in enumerate(plates_of_interest):
-            residual_x = _numpy.zeros_like(sp_const_grid); residual_y = _numpy.zeros_like(sp_const_grid); residual_z = _numpy.zeros_like(sp_const_grid)
-            if self.settings.options[case]["Slab pull torque"] and "slab_pull_torque_x" in selected_plates.columns:
-                residual_x -= selected_plates.slab_pull_torque_x.iloc[k] * sp_const_grid / self.settings.options[case]["Slab pull constant"]
-                residual_y -= selected_plates.slab_pull_torque_y.iloc[k] * sp_const_grid / self.settings.options[case]["Slab pull constant"]
-                residual_z -= selected_plates.slab_pull_torque_z.iloc[k] * sp_const_grid / self.settings.options[case]["Slab pull constant"]
-
-            # Add GPE torque
-            if self.settings.options[case]["GPE torque"] and "GPE_torque_x" in selected_plates.columns:
-                residual_x -= selected_plates.GPE_torque_x.iloc[k] * ones_grid
-                residual_y -= selected_plates.GPE_torque_y.iloc[k] * ones_grid
-                residual_z -= selected_plates.GPE_torque_z.iloc[k] * ones_grid
+        for k, _plateID in enumerate(_plateIDs):
+            if _plateID in _data.plateID.values:
+                residual_x = _numpy.zeros_like(sp_const_grid); residual_y = _numpy.zeros_like(sp_const_grid); residual_z = _numpy.zeros_like(sp_const_grid)
+                if self.settings.options[case]["Slab pull torque"] and "slab_pull_torque_x" in _data.columns:
+                    residual_x -= _data.slab_pull_torque_x.iloc[k] * sp_const_grid / self.settings.options[case]["Slab pull constant"]
+                    residual_y -= _data.slab_pull_torque_y.iloc[k] * sp_const_grid / self.settings.options[case]["Slab pull constant"]
+                    residual_z -= _data.slab_pull_torque_z.iloc[k] * sp_const_grid / self.settings.options[case]["Slab pull constant"]
+
+                # Add GPE torque
+                if self.settings.options[case]["GPE torque"] and "GPE_torque_x" in _data.columns:
+                    residual_x -= _data.GPE_torque_x.iloc[k] * ones_grid
+                    residual_y -= _data.GPE_torque_y.iloc[k] * ones_grid
+                    residual_z -= _data.GPE_torque_z.iloc[k] * ones_grid
+
+                # Compute magnitude of driving torque
+                if weight_by_area:
+                    driving_mag += _numpy.sqrt(residual_x**2 + residual_y**2 + residual_z**2) * _data.area.iloc[k] / total_area
+                else:
+                    driving_mag += _numpy.sqrt(residual_x**2 + residual_y**2 + residual_z**2) / _data.area.iloc[k]
+
+                # Add slab bend torque
+                if self.settings.options[case]["Slab bend torque"] and "slab_bend_torque_x" in _data.columns:
+                    residual_x -= _data.slab_bend_torque_x.iloc[k] * ones_grid
+                    residual_y -= _data.slab_bend_torque_y.iloc[k] * ones_grid
+                    residual_z -= _data.slab_bend_torque_z.iloc[k] * ones_grid
+
+                # Add mantle drag torque
+                if self.settings.options[case]["Mantle drag torque"] and "mantle_drag_torque_x" in _data.columns:
+                    residual_x -= _data.mantle_drag_torque_x.iloc[k] * visc_grid / self.settings.options[case]["Mantle viscosity"]
+                    residual_y -= _data.mantle_drag_torque_y.iloc[k] * visc_grid / self.settings.options[case]["Mantle viscosity"]
+                    residual_z -= _data.mantle_drag_torque_z.iloc[k] * visc_grid / self.settings.options[case]["Mantle viscosity"]
+
+                # Compute magnitude of residual
+                if weight_by_area:
+                    residual_mag += _numpy.sqrt(residual_x**2 + residual_y**2 + residual_z**2) * _data.area.iloc[k] / total_area
+                else:
+                    residual_mag += _numpy.sqrt(residual_x**2 + residual_y**2 + residual_z**2) / _data.area.iloc[k]
 
-            # Compute magnitude of driving torque
-            if weight_by_area:
-                self.driving_torque[age][case] += _numpy.sqrt(residual_x**2 + residual_y**2 + residual_z**2) * selected_plates.area.iloc[k] / total_area
-            else:
-                self.driving_torque[age][case] += _numpy.sqrt(residual_x**2 + residual_y**2 + residual_z**2) / selected_plates.area.iloc[k]
-
-            # Add slab bend torque
-            if self.settings.options[case]["Slab bend torque"] and "slab_bend_torque_x" in selected_plates.columns:
-                residual_x -= selected_plates.slab_bend_torque_x.iloc[k] * ones_grid
-                residual_y -= selected_plates.slab_bend_torque_y.iloc[k] * ones_grid
-                residual_z -= selected_plates.slab_bend_torque_z.iloc[k] * ones_grid
-
-            # Add mantle drag torque
-            if self.settings.options[case]["Mantle drag torque"] and "mantle_drag_torque_x" in selected_plates.columns:
-                residual_x -= selected_plates.mantle_drag_torque_x.iloc[k] * visc_grid / self.mech.La / self.settings.options[case]["Mantle viscosity"]
-                residual_y -= selected_plates.mantle_drag_torque_y.iloc[k] * visc_grid / self.mech.La / self.settings.options[case]["Mantle viscosity"]
-                residual_z -= selected_plates.mantle_drag_torque_z.iloc[k] * visc_grid / self.mech.La / self.settings.options[case]["Mantle viscosity"]
-
-            # Compute magnitude of residual
-            if weight_by_area:
-                self.residual_torque[age][case] += _numpy.sqrt(residual_x**2 + residual_y**2 + residual_z**2) * selected_plates.area.iloc[k] / total_area
-            else:
-                self.residual_torque[age][case] += _numpy.sqrt(residual_x**2 + residual_y**2 + residual_z**2) / selected_plates.area.iloc[k]
-
-            # Divide residual by driving torque
-            self.residual_torque_normalised[age][case] = _numpy.log10(self.residual_torque[age][case] / self.driving_torque[age][case])
+        # Divide residual by driving torque
+        residual_mag_normalised = _numpy.log10(residual_mag / driving_mag)
 
         # Find the indices of the minimum value directly using _numpy.argmin
-        self.opt_i[age][case], self.opt_j[age][case] = _numpy.unravel_index(_numpy.argmin(self.residual_torque_normalised[age][case]), self.residual_torque_normalised[age][case].shape)
-        self.opt_visc[age][case] = visc_grid[self.opt_i[age][case], self.opt_j[age][case]]
-        self.opt_sp_const[age][case] = sp_const_grid[self.opt_i[age][case], self.opt_j[age][case]]
+        opt_i, opt_j = _numpy.unravel_index(_numpy.argmin(residual_mag), residual_mag.shape)
+        opt_visc = visc_grid[opt_i, opt_j]
+        opt_sp_const = sp_const_grid[opt_i, opt_j]
 
         # Plot
         if plot == True:
-            fig, ax = plt.subplots(figsize=(15*self.constants.cm2in, 12*self.constants.cm2in))
-            im = ax.imshow(self.residual_torque_normalised[age][case], cmap="cmc.lapaz_r", vmin=-1.5, vmax=1.5)
+            fig, ax = plt.subplots(figsize=(15, 12))
+            im = ax.imshow(residual_mag_normalised, cmap="cmc.lapaz_r", vmin=-1.5, vmax=1.5)
             ax.set_yticks(_numpy.linspace(0, grid_size - 1, 5))
             ax.set_xticks(_numpy.linspace(0, grid_size - 1, 5))
             ax.set_xticklabels(["{:.2e}".format(visc) for visc in _numpy.linspace(viscosity_range[0], viscosity_range[1], 5)])
             ax.set_yticklabels(["{:.2f}".format(sp_const) for sp_const in _numpy.linspace(sp_consts.min(), sp_consts.max(), 5)])
             ax.set_xlabel("Mantle viscosity [Pa s]")
             ax.set_ylabel("Slab pull reduction factor")
-            ax.scatter(self.opt_j[age][case], self.opt_i[age][case], marker="*", facecolor="none", edgecolor="k", s=30)  # Adjust the marker style and size as needed
+            ax.scatter(opt_j, opt_i, marker="*", facecolor="none", edgecolor="k", s=30)  # Adjust the marker style and size as needed
             fig.colorbar(im, label = "Log(residual torque/driving torque)")
             plt.show()
 
         # Print results
-        print(f"Optimal coefficients for ", ", ".join(selected_plates.name.astype(str)), " plate(s), (PlateIDs: ", ", ".join(selected_plates.plateID.astype(str)), ")")
-        print("Minimum residual torque: {:.2%} of driving torque".format(10**(_numpy.amin(self.residual_torque_normalised[age][case]))))
-        print("Optimum viscosity [Pa s]: {:.2e}".format(self.opt_visc[age][case]))
-        print("Optimum Drag Coefficient [Pa s/m]: {:.2e}".format(self.opt_visc[age][case] / self.mech.La))
-        print("Optimum Slab Pull constant: {:.2%}".format(self.opt_sp_const[age][case]))
+        print(f"Optimal coefficients for ", ", ".join(_data.name.astype(str)), " plate(s), (PlateIDs: ", ", ".join(_data.plateID.astype(str)), ")")
+        print("Minimum residual torque: {:.2%} of driving torque".format(10**(_numpy.amin(residual_mag_normalised))))
+        print("Optimum viscosity [Pa s]: {:.2e}".format(opt_visc))
+        print("Optimum Drag Coefficient [Pa s/m]: {:.2e}".format(opt_visc / self.settings.mech.La))
+        print("Optimum Slab Pull constant: {:.2%}".format(opt_sp_const))
 
-        return self.opt_sp_const[age][case], self.opt_visc[age][case], self.residual_torque_normalised[age][case], (self.opt_i[age][case], self.opt_j[age][case])
+        # return self.opt_sp_const[age][case], self.opt_visc[age][case], self.residual_torque_normalised[age][case], (self.opt_i[age][case], self.opt_j[age][case])
 
     def find_slab_pull_coefficient(
             self,

plato/utils_calc.py

@@ -764,13 +764,19 @@ def sum_torque(
     centroid_position = geocentric_spherical2cartesian(plates.centroid_lat, plates.centroid_lon, constants.mean_Earth_radius_m)
 
     # Calculate the torque vector as the cross product of the Cartesian torque vector (x, y, z) with the position vector of the centroid
-    force_at_centroid = _numpy.cross(summed_torques_cartesian, centroid_position, axis=0)
+    force_at_centroid_xyz = _numpy.cross(summed_torques_cartesian, centroid_position, axis=0)
 
     # Compute force magnitude at centroid
-    plates.loc[:, f"{torque_type}_force_lat"], plates.loc[:, f"{torque_type}_force_lon"], plates.loc[:, f"{torque_type}_force_mag"], plates.loc[:, f"{torque_type}_force_azi"] = geocentric_cartesian2spherical(
-        force_at_centroid[0], force_at_centroid[1], force_at_centroid[2]
+    force_at_centroid_sph = tangent_cartesian2spherical(
+        force_at_centroid_xyz.T, plates.centroid_lat, plates.centroid_lon
     )
 
+    # Assign force components to DataFrame
+    plates.loc[:, f"{torque_type}_force_lat"] = force_at_centroid_sph[0]
+    plates.loc[:, f"{torque_type}_force_lon"] = force_at_centroid_sph[1]
+    plates.loc[:, f"{torque_type}_force_mag"] = force_at_centroid_sph[2]
+    plates.loc[:, f"{torque_type}_force_azi"] = force_at_centroid_sph[3]
+
     return plates
 
 def compute_residual_force(