Write One Block Phantom Files

sarracen/__init__.py

@@ -6,6 +6,8 @@
 from .readers.read_shamrock import read_shamrock
 from .readers.read_gasoline import read_gasoline
 
+from .writers.write_phantom import write_phantom
+
 from .sarracen_dataframe import SarracenDataFrame
 
 from .interpolate import interpolate_2d, interpolate_2d_line, \
@@ -14,4 +16,4 @@
 
 import sarracen.disc
 
-__version__ = "1.2.3"
+__version__ = "1.2.4"

sarracen/readers/read_phantom.py

@@ -22,7 +22,7 @@ def _read_fortran_block(fp, bytesize):
     return data
 
 
-def _read_capture_pattern(fp):
+def read_capture_pattern(fp):
     """ Phantom dump validation plus default real and int sizes."""
 
     start_tag = fp.read(4)  # 4-byte Fortran tag
@@ -131,7 +131,7 @@ def _read_global_header(fp, def_int_dtype, def_real_dtype):
         new_keys, new_data = _read_global_header_block(fp, dtype)
 
         keys += new_keys
-        data = np.append(data, new_data)
+        data = data + list(new_data)
 
     keys = _rename_duplicates(keys)
 
@@ -223,6 +223,13 @@ def _create_aprmass_column(df, header_vars):
     return df
 
 
+# def update_int64_header_vars(header_vars):
+#     for key, value in header_vars.items():
+#         if key in ['nparttot', 'ntypes', 'npartoftype']:
+#             header_vars[key] = np.int64(value)
+#     return header_vars
+
+
 def read_phantom(filename: str,
                  separate_types: str = 'sinks',
                  ignore_inactive: bool = True):
@@ -274,12 +281,15 @@ def read_phantom(filename: str,
     >>> sdf_gas, sdf_dust, sdf_sinks = sarracen.read_phantom('dumpfile_00000', separate_types='all')
     """
     with open(filename, 'rb') as fp:
-        def_int_dtype, def_real_dtype, iversion = _read_capture_pattern(fp)
+        def_int_dtype, def_real_dtype, iversion = read_capture_pattern(fp)
         file_identifier = _read_file_identifier(fp)
 
         header_vars = _read_global_header(fp, def_int_dtype, def_real_dtype)
+        # header_vars = update_int64_header_vars(header_vars)
         header_vars['file_identifier'] = file_identifier
         header_vars['iversion'] = iversion
+        header_vars['def_int_dtype'] = def_int_dtype
+        header_vars['def_real_dtype'] = def_real_dtype
 
         df, df_sinks = _read_array_blocks(fp, def_int_dtype, def_real_dtype)
 

sarracen/readers/read_shamrock.py

@@ -2,7 +2,7 @@
 import pandas as pd
 
 import vtk
-from vtk.util.numpy_support import vtk_to_numpy
+from vtkmodules.util.numpy_support import vtk_to_numpy
 
 from ..sarracen_dataframe import SarracenDataFrame
 

sarracen/tests/readers/test_read_phantom.py

@@ -47,6 +47,7 @@ def _create_global_header(massoftype=1e-6, massoftype_7=None,
     for i in range(8):  # loop over 8 dtypes
         file += bytearray(read_tag.tobytes())
         nvars = (i == 5) + (massoftype_7 is not None)
+        nvars = nvars + 3 # ['nparttot', 'ntypes', 'npartoftype']
         if i == 5:  # default real
             nvars = np.array([nvars], dtype='int32')
         else:
@@ -57,13 +58,19 @@ def _create_global_header(massoftype=1e-6, massoftype_7=None,
         if i == 5:  # default real
             file += bytearray(read_tag.tobytes())
             file += bytearray(map(ord, "massoftype".ljust(16)))
+            file += bytearray(map(ord, "nparttot".ljust(16)))
+            file += bytearray(map(ord, "ntypes".ljust(16)))
+            file += bytearray(map(ord, "npartoftype".ljust(16)))
             if massoftype_7 is not None:
                 file += bytearray(map(ord, "massoftype_7".ljust(16)))
             file += bytearray(read_tag.tobytes())
 
         if i == 5:
             file += bytearray(read_tag.tobytes())
             file += bytearray(np.array([massoftype], dtype=def_real))
+            file += bytearray(np.array([100], dtype=def_real))
+            file += bytearray(np.array([5], dtype=def_real))
+            file += bytearray(np.array([20], dtype=def_real))
             if massoftype_7 is not None:
                 file += bytearray(np.array([massoftype_7], dtype=def_real))
             file += bytearray(read_tag.tobytes())
@@ -171,6 +178,12 @@ def test_gas_particles_only():
         sdf = sarracen.read_phantom(fp.name, separate_types=None)
         assert sdf.params['massoftype'] == 1e-6
         assert sdf.params['mass'] == 1e-6
+        assert sdf.params['nparttot'].dtype == np.int64
+        assert sdf.params['ntypes'].dtype == np.int64
+        assert sdf.params['npartoftype'].dtype == np.int64
+        assert sdf.params['nparttot'] == 100
+        assert sdf.params['ntypes'] == 5
+        assert sdf.params['npartoftype'] == 20
         assert 'mass' not in sdf.columns
         tm.assert_series_equal(sdf['x'], pd.Series([0, 0, 0, 0, 1, 1, 1, 1]),
                                check_index=False, check_names=False,

sarracen/tests/writers/test_write_phantom.py

@@ -0,0 +1,250 @@
+import pandas as pd
+import numpy as np
+import sarracen
+import tempfile
+
+
+def _create_capture_pattern(def_int, def_real):
+    """ Construct capture pattern. """
+
+    read_tag = np.array([13], dtype='int32')
+    i1 = np.array([60769], dtype=def_int)
+    r2 = np.array([60878], dtype=def_real)
+    i2 = np.array([60878], dtype=def_int)
+    iversion = np.array([0], dtype=def_int)
+    i3 = np.array([690706], dtype=def_int)
+
+    capture_pattern = bytearray(read_tag.tobytes())
+    capture_pattern += bytearray(i1.tobytes())
+    capture_pattern += bytearray(r2.tobytes())
+    capture_pattern += bytearray(i2.tobytes())
+    capture_pattern += bytearray(iversion.tobytes())
+    capture_pattern += bytearray(i3.tobytes())
+    capture_pattern += bytearray(read_tag.tobytes())
+
+    return capture_pattern
+
+
+def _create_file_identifier():
+    """ Construct 100-character file identifier. """
+
+    read_tag = np.array([13], dtype='int32')
+    file_identifier = "Test of read_phantom".ljust(100)
+    bytes_file = bytearray(read_tag.tobytes())
+    bytes_file += bytearray(map(ord, file_identifier))
+    bytes_file += bytearray(read_tag.tobytes())
+    return bytes_file
+
+
+def _create_global_header(massoftype=1e-6, massoftype_7=None,
+                          def_int=np.int32, def_real=np.float64):
+    """ Construct global variables. Only massoftype in this example. """
+
+    read_tag = np.array([13], dtype='int32')
+    bytes_file = bytearray()
+    for i in range(8):  # loop over 8 dtypes
+        bytes_file += bytearray(read_tag.tobytes())
+        nvars = (i == 5) + (massoftype_7 is not None)
+        if i == 5:  # default real
+            nvars = np.array([nvars], dtype='int32')
+        else:
+            nvars = np.array([0], dtype='int32')
+        bytes_file += bytearray(nvars.tobytes())
+        bytes_file += bytearray(read_tag.tobytes())
+
+        if i == 5:  # default real
+            bytes_file += bytearray(read_tag.tobytes())
+            bytes_file += bytearray(map(ord, "massoftype".ljust(16)))
+            if massoftype_7 is not None:
+                bytes_file += bytearray(map(ord, "massoftype_7".ljust(16)))
+            bytes_file += bytearray(read_tag.tobytes())
+
+        if i == 5:
+            bytes_file += bytearray(read_tag.tobytes())
+            bytes_file += bytearray(np.array([massoftype], dtype=def_real))
+            if massoftype_7 is not None:
+                bytes_file += bytearray(np.array([massoftype_7], dtype=def_real))
+            bytes_file += bytearray(read_tag.tobytes())
+
+    return bytes_file
+
+
+def _create_particle_array(tag, data, dtype=np.float64):
+    read_tag = np.array([13], dtype='int32')
+    bytes_file = bytearray(read_tag.tobytes())
+    bytes_file += bytearray(map(ord, tag.ljust(16)))
+    bytes_file += bytearray(read_tag.tobytes())
+
+    bytes_file += bytearray(read_tag.tobytes())
+    bytes_file += bytearray(np.array(data, dtype=dtype).tobytes())
+    bytes_file += bytearray(read_tag.tobytes())
+    return bytes_file
+
+
+def get_one_block_phantom_file():
+    bytes_file = _create_capture_pattern(np.int32, np.float64)
+    bytes_file += _create_file_identifier()
+    bytes_file += _create_global_header()
+
+    # create 1 block for gas
+    read_tag = np.array([13], dtype='int32')
+    bytes_file += bytearray(read_tag.tobytes())
+    nblocks = np.array([1], dtype='int32')
+    bytes_file += bytearray(nblocks.tobytes())
+    bytes_file += bytearray(read_tag.tobytes())
+
+    # 8 particles storing 4 real arrays (x, y, z, h)
+    bytes_file += bytearray(read_tag.tobytes())
+    n = np.array([8], dtype='int64')
+    nums = np.array([0, 0, 0, 0, 0, 4, 0, 0], dtype='int32')
+    bytes_file += bytearray(n.tobytes())
+    bytes_file += bytearray(nums.tobytes())
+    bytes_file += bytearray(read_tag.tobytes())
+
+    # write 4 particle arrays
+    bytes_file += _create_particle_array("x", [0, 0, 0, 0, 1, 1, 1, 1])
+    bytes_file += _create_particle_array("y", [0, 0, 1, 1, 0, 0, 1, 1])
+    bytes_file += _create_particle_array("z", [0, 1, 0, 1, 0, 1, 0, 1])
+    bytes_file += _create_particle_array("h", [1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1])
+
+    return bytes_file
+
+
+def get_gas_dust_sink_particles():
+
+    bytes_file = _create_capture_pattern(np.int32, np.float64)
+    bytes_file += _create_file_identifier()
+    bytes_file += _create_global_header(massoftype_7=1e-4)
+
+    # create 1 block for gas
+    read_tag = np.array([13], dtype='int32')
+    bytes_file += bytearray(read_tag.tobytes())
+    nblocks = np.array([2], dtype='int32')
+    bytes_file += bytearray(nblocks.tobytes())
+    bytes_file += bytearray(read_tag.tobytes())
+
+    # 8 particles storing 4 real arrays (x, y, z, h)
+    bytes_file += bytearray(read_tag.tobytes())
+    n = np.array([16], dtype='int64')
+    nums = np.array([0, 1, 0, 0, 0, 4, 0, 0] , dtype='int32')
+    bytes_file += bytearray(n.tobytes())
+    bytes_file += bytearray(nums.tobytes())
+    bytes_file += bytearray(read_tag.tobytes())
+
+    bytes_file += bytearray(read_tag.tobytes())
+    n = np.array([1], dtype='int64')
+    nums = np.array([0, 0, 0, 0, 0, 7, 0, 0] , dtype='int32')
+    bytes_file += bytearray(n.tobytes())
+    bytes_file += bytearray(nums.tobytes())
+    bytes_file += bytearray(read_tag.tobytes())
+
+    # write 5 gas/dust particle arrays
+    bytes_file += _create_particle_array("itype", [1, 1, 1, 1, 1, 1, 1, 1,
+                                         7, 7, 7, 7, 7, 7, 7, 7], np.int8)
+    bytes_file += _create_particle_array("x", [0, 0, 0, 0, 1, 1, 1, 1,
+                                               0.5, 0.5, 0.5, 0.5, 1.5, 1.5, 1.5, 1.5])
+    bytes_file += _create_particle_array("y", [0, 0, 1, 1, 0, 0, 1, 1,
+                                               0.5, 0.5, 1.5, 1.5, 0.5, 0.5, 1.5, 1.5])
+    bytes_file += _create_particle_array("z", [0, 1, 0, 1, 0, 1, 0, 1,
+                                               0.5, 1.5, 0.5, 1.5, 0.5, 1.5, 0.5, 1.5])
+    bytes_file += _create_particle_array("h", [1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1,
+                                               1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1])
+
+    # write 7 sink particle arrays
+    bytes_file += _create_particle_array("x", [0.000305])
+    bytes_file += _create_particle_array("y", [-0.035809])
+    bytes_file += _create_particle_array("z", [-0.000035])
+    bytes_file += _create_particle_array("h", [1.0])
+    bytes_file += _create_particle_array("spinx", [-3.911744e-8])
+    bytes_file += _create_particle_array("spiny", [-1.326062e-8])
+    bytes_file += _create_particle_array("spinz", [0.00058])
+    return bytes_file
+
+def get_gas_sink_particles():
+
+    bytes_file = _create_capture_pattern(np.int32, np.float64)
+    bytes_file += _create_file_identifier()
+    bytes_file += _create_global_header()
+
+    # create 1 block for gas
+    read_tag = np.array([13], dtype='int32')
+    bytes_file += bytearray(read_tag.tobytes())
+    nblocks = np.array([2], dtype='int32')
+    bytes_file += bytearray(nblocks.tobytes())
+    bytes_file += bytearray(read_tag.tobytes())
+
+    # 8 particles storing 4 real arrays (x, y, z, h)
+    bytes_file += bytearray(read_tag.tobytes())
+    n = np.array([8], dtype='int64')
+    nums = np.array([0, 0, 0, 0, 0, 4, 0, 0] , dtype='int32')
+    bytes_file += bytearray(n.tobytes())
+    bytes_file += bytearray(nums.tobytes())
+    bytes_file += bytearray(read_tag.tobytes())
+
+    bytes_file += bytearray(read_tag.tobytes())
+    n = np.array([1], dtype='int64')
+    nums = np.array([0, 0, 0, 0, 0, 7, 0, 0] , dtype='int32')
+    bytes_file += bytearray(n.tobytes())
+    bytes_file += bytearray(nums.tobytes())
+    bytes_file += bytearray(read_tag.tobytes())
+
+    # write 4 gas particle arrays
+    bytes_file += _create_particle_array("x", [0, 0, 0, 0, 1, 1, 1, 1])
+    bytes_file += _create_particle_array("y", [0, 0, 1, 1, 0, 0, 1, 1])
+    bytes_file += _create_particle_array("z", [0, 1, 0, 1, 0, 1, 0, 1])
+    bytes_file += _create_particle_array("h", [1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1])
+
+    # write 7 sink particle arrays
+    bytes_file += _create_particle_array("x", [0.000305])
+    bytes_file += _create_particle_array("y", [-0.035809])
+    bytes_file += _create_particle_array("z", [-0.000035])
+    bytes_file += _create_particle_array("h", [1.0])
+    bytes_file += _create_particle_array("spinx", [-3.911744e-8])
+    bytes_file += _create_particle_array("spiny", [-1.326062e-8])
+    bytes_file += _create_particle_array("spinz", [0.00058])
+
+    return bytes_file
+
+
+def get_file_content(file_path):
+    with open(file_path, 'rb') as file:
+        return file.read()
+
+
+def test_sink_particles_gas_and_dust(): #PASSES
+    with tempfile.NamedTemporaryFile() as fp:
+        fp.write(get_gas_dust_sink_particles())
+        fp.seek(0)
+
+        test_sdfs = sarracen.read_phantom(fp.name)
+        phantom_file = sarracen.write_phantom("wp", test_sdfs[0], test_sdfs[1])
+        test_sdfs_from_new_file = sarracen.read_phantom(phantom_file.name)
+        pd.testing.assert_frame_equal(test_sdfs_from_new_file[0], test_sdfs[0])
+        pd.testing.assert_frame_equal(test_sdfs_from_new_file[1], test_sdfs[1])
+        #
+        # original_content = get_file_content(fp.name)
+        # new_content = get_file_content(phantom_file.name)
+        #
+        # assert original_content == new_content
+
+def test_sink_particles(): #PASSES
+    with tempfile.NamedTemporaryFile() as fp:
+        fp.write(get_gas_sink_particles())
+        fp.seek(0)
+
+        test_sdfs = sarracen.read_phantom(fp.name)
+        phantom_file = sarracen.write_phantom("wp", test_sdfs[0], test_sdfs[1])
+        test_sdfs_from_new_file = sarracen.read_phantom(phantom_file.name)
+        pd.testing.assert_frame_equal(test_sdfs_from_new_file[0], test_sdfs[0])
+        pd.testing.assert_frame_equal(test_sdfs_from_new_file[1], test_sdfs[1])
+
+
+def test_write_phantom_one_block(): #PASSES
+    with tempfile.NamedTemporaryFile() as fp:
+        fp.write(get_one_block_phantom_file())
+        fp.seek(0)
+
+        test_sdf = sarracen.read_phantom(fp.name)
+        phantom_file = sarracen.write_phantom("wp", test_sdf)
+        test_sdf_from_new_file = sarracen.read_phantom(phantom_file.name)
+        pd.testing.assert_frame_equal(test_sdf, test_sdf_from_new_file)