Sum velocity-binned vpkt spectra across all ranks

artistools/misc.py

@@ -258,6 +258,27 @@ def get_vpkt_config(modelpath: Path | str) -> dict[str, t.Any]:
             int(x) for x in vpkt_txt.readline().split()
         )
 
+        # read the next line
+        vpackets_total_active_spectral_region = vpkt_txt.readline().split()
+        vpkt_config["vpackets total active spectral region"] = vpackets_total_active_spectral_region
+
+        optically_thick_cells = vpkt_txt.readline().split()
+        vpkt_config["optically thick cells"] = optically_thick_cells
+
+        optical_depths = vpkt_txt.readline().split()
+        vpkt_config["optical depths"] = optical_depths
+
+        velocity_grid_map = vpkt_txt.readline().split()
+        vpkt_config["velocity grid map"] = velocity_grid_map
+
+        velocity_grid_map_active = vpkt_txt.readline().split()
+        vpkt_config["velocity grid map active"] = velocity_grid_map_active
+
+        velocity_grid_map_regions = vpkt_txt.readline().split()
+        vpkt_config["Number Of Velocity Maps"] = velocity_grid_map_regions[0]
+
+        vpkt_config["Velocity Map Regions"] = velocity_grid_map_regions[1:]
+
     return vpkt_config
 
 

artistools/spectra/__init__.py

@@ -16,6 +16,7 @@
 from artistools.spectra.spectra import get_vspecpol_data as get_vspecpol_data
 from artistools.spectra.spectra import get_vspecpol_spectrum as get_vspecpol_spectrum
 from artistools.spectra.spectra import make_averaged_vspecfiles as make_averaged_vspecfiles
+from artistools.spectra.spectra import make_virtual_grid_summed_file as make_virtual_grid_summed_file
 from artistools.spectra.spectra import make_virtual_spectra_summed_file as make_virtual_spectra_summed_file
 from artistools.spectra.spectra import print_integrated_flux as print_integrated_flux
 from artistools.spectra.spectra import read_spec_res as read_spec_res

artistools/spectra/plotspectra.py

@@ -1341,6 +1341,10 @@ def addargs(parser) -> None:
         "--makevspecpol", action="store_true", help="Make file summing the virtual packet spectra from all ranks"
     )
 
+    parser.add_argument(
+        "--makevspecgrid", action="store_true", help="Make file summing the virtual packet grid from all ranks"
+    )
+
     # To get better statistics for polarisation use multiple runs of the same simulation. This will then average the
     # files produced by makevspecpol for all simulations.
     parser.add_argument(
@@ -1473,6 +1477,10 @@ def main(args: argparse.Namespace | None = None, argsraw: Sequence[str] | None =
         atspectra.make_virtual_spectra_summed_file(args.specpath[0])
         return
 
+    if args.makevspecgrid:
+        atspectra.make_virtual_grid_summed_file(args.specpath[0])
+        return
+
     if args.averagevspecpolfiles:
         atspectra.make_averaged_vspecfiles(args)
         return

artistools/spectra/spectra.py

@@ -558,6 +558,83 @@ def make_virtual_spectra_summed_file(modelpath: Path | str) -> None:
         print(f"Saved {outfile}")
 
 
+def make_virtual_grid_summed_file(modelpath: Path | str) -> None:
+    print("Generating all-rank summed vspec Grid files..")
+    nprocs = get_nprocs(modelpath)
+    print("nprocs", nprocs)
+    vpktconfig = get_vpkt_config(modelpath)
+
+    nvirtual_spectra = int(vpktconfig["nobsdirections"])
+    nvmaps = int(vpktconfig["Number Of Velocity Maps"])
+
+    # Read the first file to get the dimensions
+    vpkt_grid_files = [f"vpkt_grid_{mpirank:04d}.out" for mpirank in range(nprocs)]
+    vpkt_grid_data = pd.read_csv(vpkt_grid_files[0], sep=" ", header=None)
+    VGRID_NY = len(vpkt_grid_data[0].unique())
+    VGRID_NZ = len(vpkt_grid_data[1].unique())
+
+    print(f"VGRID_NY: {VGRID_NY}, VGRID_NZ: {VGRID_NZ}")
+
+    total_grid_points = VGRID_NY * VGRID_NZ  # Total grid points in one layer
+    rows_per_obsdir = total_grid_points * nvmaps  # Rows for one observer direction
+
+    # Output filename
+    output_filename_template = "Vpkt_grid_total_{}.txt"
+
+    # Initialize variables for summation
+    summed_data = None
+    velocity_columns = None  # velocity separately
+
+    for filename in vpkt_grid_files:
+        data = np.loadtxt(filename)
+        print(f"Processing {filename}")
+
+        # Validate row count
+        expected_rows = rows_per_obsdir * nvirtual_spectra
+
+        if data.shape[0] != expected_rows:
+            print(f"Unexpected number of rows in {filename}. Expected {expected_rows}, got {data.shape[0]}.")
+
+        # Separate velocity and other columns
+        if velocity_columns is None:
+            velocity_columns = data[:, :2]  # Assuming the first two columns are velocity
+        else:
+            print(f"Velocity columns in {filename} do not match previous files.")
+
+        # Sum the remaining columns
+        if summed_data is None:
+            summed_data = data[:, 2:]  # Exclude velocity columns from summation
+        else:
+            summed_data += data[:, 2:]
+    print("Files successfully processed.")
+    # Combine velocity and summed data
+    valid_data = tuple(array for array in (velocity_columns, summed_data) if array is not None)
+    # Perform the horizontal stack
+    final_data = np.hstack(valid_data)
+
+    # Ensure final data has exactly 5 columns
+    if final_data.shape[1] > 5:
+        print("Data has more than 5 columns. Should only have N1, N2, I, Q, U.")
+    if final_data.shape[1] < 5:
+        print("Data has less than 5 columns. Should only have N1, N2, I, Q, U.")
+
+    # Split combined data by observer direction and write to output files
+    for obsdir in range(nvirtual_spectra):
+        start_idx = obsdir * rows_per_obsdir
+        end_idx = start_idx + rows_per_obsdir
+
+        # Extract data for this observer direction
+        obs_data = final_data[start_idx:end_idx, :]
+
+        # Write to output file
+        output_path = Path(output_filename_template.format(obsdir))
+        with output_path.open("w", encoding="utf-8") as f:  # Specify encoding
+            for row in obs_data:
+                f.write(", ".join(map(str, row)) + "\n")
+
+    print("Velocity Grid successfully Written.")
+
+
 def make_averaged_vspecfiles(args: argparse.Namespace) -> None:
     filenames = [
         vspecfile.name