diff --git a/artistools/misc.py b/artistools/misc.py index 7ea07442c..01bcc40b7 100644 --- a/artistools/misc.py +++ b/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 diff --git a/artistools/spectra/__init__.py b/artistools/spectra/__init__.py index 3452d0bab..d08647320 100644 --- a/artistools/spectra/__init__.py +++ b/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 diff --git a/artistools/spectra/plotspectra.py b/artistools/spectra/plotspectra.py index d57da7eb3..9af60abb2 100755 --- a/artistools/spectra/plotspectra.py +++ b/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 diff --git a/artistools/spectra/spectra.py b/artistools/spectra/spectra.py index a390a2e54..2c0d485d6 100644 --- a/artistools/spectra/spectra.py +++ b/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