swap orient_stars with orient_component, takes a string now.

valid options are star(s), gas, young_star(s), and cold_gas.
extraction now defaults to cold_gas w/i 0.1 rvir

Author: agurvich

src/abg_python/galaxy/cosmoExtractor.py

@@ -32,7 +32,7 @@ def getThetasTaitBryan(angMom):
     ## RETURNS DEGREES
     return theta_TB,phi_TB
 
-def offsetRotateSnapshot(snap,scom,vscom,theta_TB,phi_TB,orient_stars):
+def offsetRotateSnapshot(snap,scom,vscom,theta_TB,phi_TB,orient_component):
 
     if 'overwritten' in snap and snap['overwritten']:
         if (snap['theta_TB'] == theta_TB and 
@@ -61,7 +61,7 @@ def offsetRotateSnapshot(snap,scom,vscom,theta_TB,phi_TB,orient_stars):
         'theta_TB':theta_TB,
         'phi_TB':phi_TB,
         'overwritten':1,
-        'orient_stars':orient_stars}
+        'orient_component':orient_component}
 
     ## if the snap is already overwritten this info is in there
     if 'scom' not in snap:
@@ -102,7 +102,7 @@ def extractDiskFromSnapdicts(
     orient_radius,
     scom=None,
     dark_snap=None,
-    orient_stars=0,
+    orient_component='cold_gas',
     force_theta_TB=None,
     force_phi_TB=None,
     loud=True):
@@ -133,7 +133,7 @@ def extractDiskFromSnapdicts(
         snap,
         orient_radius,
         scom=scom,
-        orient_stars=orient_stars,
+        orient_component=orient_component,
         force_theta_TB=force_theta_TB,
         force_phi_TB=force_phi_TB)
     if loud: print("Done.")
@@ -144,7 +144,7 @@ def extractDiskFromSnapdicts(
             this_snap,
             scom,vscom,
             theta_TB,phi_TB,
-            orient_stars)
+            orient_component)
 
     ## dictionary to add to extracted snapshot
     add_to_dict = {
@@ -166,10 +166,10 @@ def extractDiskFromSnapdicts(
 
 def orientDiskFromSnapdicts(
     star_snap,
-    snap,
+    gas_snap,
     radius,
     scom,
-    orient_stars=0,
+    orient_component='cold_gas',
     force_theta_TB=None,
     force_phi_TB=None):
     """ Takes arrays from a snapshot and returns orientation.
@@ -179,23 +179,27 @@ def orientDiskFromSnapdicts(
             radius - radius to extract particles from
     """
 
-    if orient_stars:
-        if star_snap is None:
-            raise ValueError("Can't orient on stars if stars are not passed")
-
-        these_rs = star_snap['Coordinates']
-        these_vs = star_snap['Velocities']
-        these_masses = star_snap['Masses']
-    else:
-        if snap is None:
-            raise ValueError("Can't orient on gas if gas is not passed")
-
-        these_rs = snap['Coordinates']
-        these_vs = snap['Velocities']
-        these_masses = snap['Masses']
-
-    mask = sphericalVolumeMask(
-            these_rs,radius,scom)
+    if 'star' in orient_component: which_snap = star_snap
+    elif 'gas' in orient_component: which_snap = gas_snap
+    else: raise KeyError(
+        f"Can't understand orient_component:{orient_component}")
+
+    these_rs = which_snap['Coordinates']
+    these_vs = which_snap['Velocities']
+    these_masses = which_snap['Masses']
+
+    mask = sphericalVolumeMask(these_rs,radius,scom)
+
+    if '_' in orient_component:
+        ## hardcode in my definition of "cold", 1e3 K
+        if orient_component == 'cold_gas':
+            mask = np.logical_and(mask,which_snap['Temperature']<1e3)
+        ## hardcode in my definition of "young", 25 Myr
+        elif orient_component in ['young_star','young_stars']:
+            mask = np.logical_and(mask,which_snap['AgeGyr']<0.025)
+        ## catch for everything else
+        else: raise NotImplementedError(
+            f"Orient component {orient_component} is not implemented.")
 
     if not np.sum(mask):
         print(scom,radius)

src/abg_python/galaxy/gal_utils.py

@@ -809,7 +809,7 @@ def saveHeaderToCatalog(
     def extractMainHalo(
         self,
         save_meta=False, 
-        orient_stars=True, ## which angular momentum axis to orient along
+        orient_component='cold_gas', ## which angular momentum axis to orient along
         overwrite_full_snaps_with_rotated_versions=False,
         free_mem=True, ## delete the full snapshot from RAM
         extract_DM=True, ## do we want the DM particles? 
@@ -828,15 +828,12 @@ def extractMainHalo(
         force_phi_TB = None -- force orientation
         """
 
-        if orient_stars:
-            group_name = 'star_extract'
-        else:
-            group_name = 'gas_extract'
+        group_name = f"{orient_component}_extract"
 
         @metadata_cache(
             group_name,
             ['sub_radius',
-            'orient_stars',
+            'orient_component',
             'theta_TB',
             'phi_TB',
             'rvir',
@@ -847,7 +844,7 @@ def extractMainHalo(
             loud=loud)
         def extract_halo_inner(
             self,
-            orient_stars=True,
+            orient_component='cold_gas',
             radius=None,
             use_saved_subsnapshots=True,
             force=False,
@@ -863,19 +860,20 @@ def extract_halo_inner(
                     self.star_snap['Masses'],
                     n=4)
 
-            if radius is None:
-                radius = self.rvir ## radius of the sub-snapshot
+            ## calculate the stellar half-mass radius to have--
+            ##  some analyses require 5rstar_half
+            if self.rstar_half is None: self.get_rstar_half(
+                save_meta=save_meta,
+                force_from_file=True,
+                loud=False,
+                within_radius=0.2*self.rvir)
 
-                ## manually calcualte rstar half using the star particles
-                ##  rather than relying on the output of AHF
-                if self.rstar_half is None: self.get_rstar_half(
-                    save_meta=save_meta,
-                    force_from_file=True,
-                    loud=False)
+            if radius is None: radius = self.rvir ## radius of the sub-snapshot 
 
-                ## radius to calculate angular momentum
-                ##  to orient on 
-                orient_radius = 5*self.rstar_half 
+            ## radius to calculate angular momentum to orient on
+            #orient_radius = 5*self.rstar_half 
+            ## switched on 1/16/2023 for paper 3, also orient on cold_gas
+            orient_radius =  0.1 * self.rvir 
 
             ## if this is not the first time extract_halo_inner has been 
             ##  called these properties may or may not exist
@@ -1005,7 +1003,7 @@ def extract_halo_inner(
                 orient_radius, ## radius to orient on
                 scom=self.scom,
                 dark_snap=which_dark_snap, ## dark_snap = None will ignore dark matter particles
-                orient_stars=orient_stars,
+                orient_component=orient_component,
                 force_theta_TB=force_theta_TB,
                 force_phi_TB=force_phi_TB,
                 loud=loud)
@@ -1022,7 +1020,7 @@ def extract_halo_inner(
             self.sub_radius = orient_radius
             ## denote whether the most recent extraction was 
             ##  oriented on stars or gas
-            self.orient_stars = orient_stars
+            self.orient_component = orient_component
 
             if compute_stellar_hsml and 'SmoothingLength' not in self.sub_star_snap: 
                 already_saved=False
@@ -1042,7 +1040,7 @@ def extract_halo_inner(
 
 
             return (self.sub_radius,
-                self.orient_stars,
+                self.orient_component,
                 self.sub_snap['theta_TB'],
                 self.sub_snap['phi_TB'],
                 self.rvir,
@@ -1051,7 +1049,7 @@ def extract_halo_inner(
 
         return_value = extract_halo_inner(
             self,
-            orient_stars=orient_stars,
+            orient_component=orient_component,
             **kwargs)
 
         ## this should happen by default when you do an extraction but if you are 
@@ -1103,10 +1101,10 @@ def get_rstar_half(self,
             loud=loud,
             assert_cached=assert_cached,
             force_from_file=force_from_file)
-        def compute_rstar_half(self):
+        def compute_rstar_half(self,within_radius=None):
             self.load_stars()
-            return self.calculate_half_mass_radius(),
-        return compute_rstar_half(self)
+            return self.calculate_half_mass_radius(within_radius=within_radius),
+        return compute_rstar_half(self,**kwargs)
 
     def get_HSML(
         self,
@@ -1266,10 +1264,10 @@ def overwrite_full_snaps_with_rotated_versions(self,extract_DM):
             snaps += [self.dark_snap]
 
         ## get the extraction parameters
-        theta_TB,phi_TB,scom,vscom,orient_stars = (
+        theta_TB,phi_TB,scom,vscom,orient_component = (
             self.sub_snap['theta_TB'],self.sub_snap['phi_TB'],
             self.sub_snap['scom'],self.sub_snap['vscom'],
-            self.sub_snap['orient_stars'])
+            self.sub_snap['orient_component'])
 
         ## if snaps doesn't have dark snap then zipped will stop at [0,4]
         for ptype,snap in zip([0,4,1],snaps):
@@ -1279,7 +1277,7 @@ def overwrite_full_snaps_with_rotated_versions(self,extract_DM):
                     snap,
                     scom,vscom,
                     theta_TB,phi_TB,
-                    orient_stars)
+                    orient_component)
 
                 ## snapdict holds ptype -> "snap"/"star_snap"/"dark_snap"
                 setattr(self,snap_dict[ptype],snap)