From 173de1151c5ad2c06c984847ec653f3347a418d8 Mon Sep 17 00:00:00 2001
From: Noe Brucy <noe.brucy@cea.fr>
Date: Fri, 5 May 2023 19:15:44 +0200
Subject: [PATCH] add cartesian analysis

---
 galsec.py | 79 ++++++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 78 insertions(+), 1 deletion(-)

diff --git a/galsec.py b/galsec.py
index a9d83a2..37b683b 100644
--- a/galsec.py
+++ b/galsec.py
@@ -90,7 +90,7 @@ def aggregate(
     grouped_data: QTable,
     weight_field: str = "mass",
     extensive_fields: str = ["mass", "ek"],
-    weighted_fields: list = ["velr", "velphi", "velz"],
+    weighted_fields: list = ["velr", "velphi", "velx", "vely", "velz"],
 ) -> QTable:
     """Aggregate wisely from grouped data
 
@@ -281,6 +281,8 @@ class Galsec:
                 dset["position"], dset["velocity"]
             )  # azimuthal velocity
             dset["velr"] = vect_r(dset["position"], dset["velocity"])  # radial velocity
+            dset["velx"] = dset["velocity"][:, 0]  # x velocity (TODO this is stupid)
+            dset["vely"] = dset["velocity"][:, 1]  # y velocity
             dset["velz"] = dset["velocity"][:, 2]  # vertical velocity
             dset["ek"] = (dset["mass"] * np.sum(dset["velocity"] ** 2, axis=1)).to(
                 u.erg
@@ -316,6 +318,26 @@ class Galsec:
             delta_phi = (delta_l / self.data[fluid]["r_bin"]) * u.rad
             phi_bin = (np.trunc(self.data[fluid]["phi"] / delta_phi) + 0.5) * delta_phi
             self.data[fluid]["phi_bin"] = phi_bin
+            
+    def cartesian_binning(self, delta_x: Quantity[u.kpc], 
+                          delta_y: Quantity[u.kpc],
+                          delta_z: Quantity[u.kpc]):
+        """Add cartesian bin information to the dataset
+
+        Parameters
+        ----------
+        delta_x : Quantity[u.kpc]
+            spacing between two x bins
+        delta_y : Quantity[u.kpc]
+             spacing between two y bins
+        """
+        for fluid in self.fluids:
+            x_bin = np.trunc(self.data[fluid]["pos"][:, 0] / delta_x)
+            y_bin = np.trunc(self.data[fluid]["pos"][:, 0] / delta_y)
+            z_bin = np.trunc(self.data[fluid]["pos"][:, 0] / delta_z)
+            self.data[fluid]["x_bin"] = x_bin
+            self.data[fluid]["y_bin"] = y_bin
+            self.data[fluid]["z_bin"] = z_bin
 
     def sector_analysis(
         self,
@@ -376,6 +398,61 @@ class Galsec:
             self.sectors["stars"]["sfr"][i] = get_sfr(group, self.time)
 
         self.sectors["stars"]["sfr"][i] = get_sfr(group, self.time)
+        
+    def cartesian_analysis(
+        self,
+        delta_x: Quantity[u.kpc] = u.kpc,
+        delta_y: Quantity[u.kpc] = u.kpc,
+        delta_z: Quantity[u.kpc] = u.kpc,
+        zmin: Quantity[u.kpc] = -0.5 * u.kpc,
+        zmax: Quantity[u.kpc] = 0.5 * u.kpc,
+    ):
+        """Compute the aggration of quantities in sectors bins
+
+        Parameters
+        ----------
+        delta_x : Quantity[u.kpc]
+            spacing between two x bins
+        delta_y : Quantity[u.kpc]
+             spacing between two y bins
+        """
+
+        self.cartesian_binning(delta_x, delta_y)
+        self.grouped_data = {}
+        self.sectors = {}
+
+        for fluid in self.fluids:
+            if fluid == "gas":
+                extensive_fields = ["mass", "ek", "volume"]
+            else:
+                extensive_fields = ["mass", "ek"]
+            filtered_data = self.data[fluid][
+                (self.data[fluid]["position"][:, 2] > zmin)
+                & (self.data[fluid]["position"][:, 2] < zmax)
+            ]
+            self.grouped_data[fluid] = filtered_data.group_by(["x_bin", "y_bin", "z_bin"])
+            self.sectors[fluid] = hstack(
+                [
+                    self.grouped_data[fluid]["x_bin", "y_bin", "z_bin"].groups.aggregate(
+                        np.fmin
+                    ),
+                    aggregate(
+                        self.grouped_data[fluid], extensive_fields=extensive_fields
+                    ),
+                ]
+            )
+            self.sectors[fluid].rename_column("x_bin", "x")
+            self.sectors[fluid].rename_column("y_bin", "y")
+            self.sectors[fluid].rename_column("z_bin", "z")
+
+        self.sectors["stars"]["sfr"] = (
+            np.zeros(len(self.sectors["stars"]["mass"])) * u.Msun / u.year
+        )
+        for i, group in enumerate(self.grouped_data["stars"].groups):
+            self.sectors["stars"]["sfr"][i] = get_sfr(group, self.time)
+
+        self.sectors["stars"]["sfr"][i] = get_sfr(group, self.time)
+
 
     def ring_analysis(
         self,