dtype updated to match the latest pkdgrav3 fof

data structure. dtype can also be provided to the reader.
sambit-giri · Jan 18, 2024 · 147cea0 · 147cea0
1 parent 37a7413
commit 147cea0
Showing 1 changed file with 27 additions and 24 deletions.
diff --git a/src/tools21cm/nbody_pkdgrav.py b/src/tools21cm/nbody_pkdgrav.py
@@ -81,7 +81,7 @@ def __init__(self, box_len, nGrid,
             Omega_m=0.31, rho_c=2.77536627e11, verbose=True):
         super().__init__(box_len, nGrid, Omega_m, rho_c, verbose)
 
-    def read_fof_data(self, filename, z=None):
+    def read_fof_data(self, filename, z=None, dtype=None):
         '''
         Read the FOF data.
 
@@ -108,26 +108,27 @@ def read_fof_data(self, filename, z=None):
         self.dMassFac = dMassFac
         self.dVelFac  = dVelFac
 
-        dtype = np.dtype([
-            ('rPot', '<f4', (3,)),
-            ('minPot', '<f4'),
-            ('rcen', '<f4', (3,)),
-            ('rcom', '<f4', (3,)),
-            ('vcom', '<f4', (3,)),
-            ('angular', '<f4', (3,)),
-            ('inertia', '<f4', (6,)),
-            ('sigma', '<f4'),
-            ('rMax', '<f4'),
-            ('fMass', '<f4'),
-            ('fEnvironDensity0', '<f4'),
-            ('fEnvironDensity1', '<f4'),
-            ('rHalf', '<f4'),
-            ('nBH', '<i4'),
-            ('nStar', '<i4'),
-            ('nGas', '<i4'),
-            ('nDM', '<i4'),
-            ('iGlobalGid', '<u8')
-        ])
+        if dtype is None:
+            dtype = np.dtype([
+                ('rPot', '<f4', (3,)),
+                ('minPot', '<f4'),
+                ('rcen', '<f4', (3,)),
+                ('rcom', '<f4', (3,)),
+                ('vcom', '<f4', (3,)),
+                ('angular', '<f4', (3,)),
+                ('inertia', '<f4', (6,)),
+                ('sigma', '<f4'),
+                ('rMax', '<f4'),
+                ('fMass', '<f4'),
+                ('fEnvironDensity0', '<f4'),
+                ('fEnvironDensity1', '<f4'),
+                ('rHalf', '<f4'),
+                # ('nBH', '<i4'),
+                # ('nStar', '<i4'),
+                # ('nGas', '<i4'),
+                # ('nDM', '<i4'),
+                # ('iGlobalGid', '<u8')
+            ])
 
         data = np.fromfile(filename, dtype=dtype)
         self.fof_data_dtype = dtype
@@ -159,14 +160,16 @@ def array_fof_data(self, dtype=float):
                 vel = [dVelFac * g['vcom'][j] for j in range(3)]
                 fMass = dMassFac * g['fMass']
                 if dtype in ['str','string',str]: 
-                    line = f"{fMass} {pos[0]:20.14f} {pos[1]:20.14f} {pos[2]:20.14f} {g['iGlobalGid']}"
+                    # line = f"{fMass} {pos[0]:20.14f} {pos[1]:20.14f} {pos[2]:20.14f} {g['iGlobalGid']}"
+                    line = f"{fMass} {pos[0]:20.14f} {pos[1]:20.14f} {pos[2]:20.14f}"
                 else:
-                    line = np.array([fMass,pos[0],pos[1],pos[2],g['iGlobalGid']]).astype(dtype)
+                    # line = np.array([fMass,pos[0],pos[1],pos[2],g['iGlobalGid']]).astype(dtype)
+                    line = np.array([fMass,pos[0],pos[1],pos[2]]).astype(dtype)
                 # print(line)
                 str_data.append(line)
             str_data = np.array(str_data)
             if self.verbose:
-                print('Returned array contains mass, pos_x, pos_y, pos_z, particle_ID.')
+                print('Returned array contains mass, pos_x, pos_y, pos_z.')
             return str_data 
         else:
             print('Data unread. Use the read_fof_data attribute.')