134 lines
No EOL
5.2 KiB
Python
134 lines
No EOL
5.2 KiB
Python
from pysbmy.density import mesh_to_mesh
|
|
from pysbmy.field import Field, read_field
|
|
import numpy as np
|
|
import os
|
|
|
|
|
|
def field_to_field(
|
|
input_file:str,
|
|
output_file:str,
|
|
output_size:int|tuple[int,int,int]|list[int],
|
|
output_L:float|tuple[float,float,float]|list[float],
|
|
output_dpm:float|tuple[float,float,float]|list[float],
|
|
output_corner:tuple[float,float,float]|list[float],
|
|
boundary_conditions:int,
|
|
):
|
|
|
|
|
|
### Make sure all inputs are valid
|
|
if output_L is not None and output_dpm is not None:
|
|
raise ValueError("Either output_L or output_dpm can be specified, not both.")
|
|
|
|
if isinstance(output_size, int):
|
|
output_size = (output_size, output_size, output_size) # N0, N1, N2
|
|
elif len(output_size) == 1:
|
|
output_size = (output_size[0], output_size[0], output_size[0])
|
|
|
|
if output_L is not None:
|
|
if isinstance(output_L, float):
|
|
output_L = (output_L, output_L, output_L)
|
|
elif len(output_L) == 1:
|
|
output_L = (output_L[0], output_L[0], output_L[0])
|
|
|
|
if output_dpm is None:
|
|
if output_L is None:
|
|
output_dpm = (-1, -1, -1)
|
|
else:
|
|
output_dpm = (output_L[0] / output_size[0], output_L[1] / output_size[1], output_L[2] / output_size[2])
|
|
elif isinstance(output_dpm, float):
|
|
output_dpm = (output_dpm, output_dpm, output_dpm) # d0, d1, d2
|
|
elif len(output_dpm) == 1:
|
|
output_dpm = (output_dpm[0], output_dpm[0], output_dpm[0])
|
|
|
|
if not os.path.exists(input_file):
|
|
raise FileNotFoundError(f"Input file {input_file} does not exist.")
|
|
|
|
|
|
|
|
# Read the input field
|
|
print(f"Reading input field from {input_file}")
|
|
input_field = read_field(input_file)
|
|
|
|
if input_field.rank != 1 or input_field.data.ndim != 3:
|
|
raise NotImplementedError("Only 3D scalar fields are supported for now.")
|
|
|
|
L0 = input_field.L0
|
|
L1 = input_field.L1
|
|
L2 = input_field.L2
|
|
N0 = input_field.N0
|
|
N1 = input_field.N1
|
|
N2 = input_field.N2
|
|
corner0 = input_field.corner0
|
|
corner1 = input_field.corner1
|
|
corner2 = input_field.corner2
|
|
|
|
d0_in = L0 / N0
|
|
d1_in = L1 / N1
|
|
d2_in = L2 / N2
|
|
|
|
d0_out = output_dpm[0] if output_dpm[0] > 0 else d0_in
|
|
d1_out = output_dpm[1] if output_dpm[1] > 0 else d1_in
|
|
d2_out = output_dpm[2] if output_dpm[2] > 0 else d2_in
|
|
|
|
offset_out_x = (output_corner[0] - corner0)
|
|
offset_out_y = (output_corner[1] - corner1)
|
|
offset_out_z = (output_corner[2] - corner2)
|
|
|
|
print("-----------------------------------------")
|
|
print(f"Input field size: {N0} x {N1} x {N2}")
|
|
print(f"Output field size: {output_size[0]} x {output_size[1]} x {output_size[2]}")
|
|
print(f"Input field dpm: {d0_in:.3f} x {d1_in:.3f} x {d2_in:.3f}")
|
|
print(f"Output field dpm: {d0_out:.3f} x {d1_out:.3f} x {d2_out:.3f}")
|
|
print(f"Input field corner: ({corner0:.1f}, {corner1:.1f}, {corner2:.1f})")
|
|
print(f"Output field corner: ({output_corner[0]:.1f}, {output_corner[1]:.1f}, {output_corner[2]:.1f})")
|
|
print(f"Boundary conditions: {'periodic' if boundary_conditions == 1 else 'non-periodic'}")
|
|
print("-----------------------------------------")
|
|
|
|
input_grid = input_field.data
|
|
|
|
output_grid = np.zeros(output_size, dtype=input_grid.dtype)
|
|
|
|
|
|
# Mesh to mesh interpolation
|
|
print("Interpolating field...")
|
|
mesh_to_mesh(input_grid, output_grid, d0_in, d1_in, d2_in, d0_out, d1_out, d2_out, offset_out_x, offset_out_y, offset_out_z, boundary_conditions)
|
|
|
|
# Create the output field
|
|
output_field = Field(
|
|
L0=output_size[0] * d0_out,
|
|
L1=output_size[1] * d1_out,
|
|
L2=output_size[2] * d2_out,
|
|
corner0=output_corner[0],
|
|
corner1=output_corner[1],
|
|
corner2=output_corner[2],
|
|
rank=input_field.rank,
|
|
N0=output_size[0],
|
|
N1=output_size[1],
|
|
N2=output_size[2],
|
|
time=input_field.time,
|
|
data=output_grid
|
|
)
|
|
|
|
# Write the output field
|
|
print(f"Writing output field to {output_file}")
|
|
output_field.write(output_file)
|
|
print("Done.")
|
|
|
|
|
|
def console_main():
|
|
import argparse
|
|
parser = argparse.ArgumentParser(description="Convert a field from one size to another.")
|
|
parser.add_argument("-i","--input_file", type=str, help="Input field file")
|
|
parser.add_argument("-o","--output_file", type=str, help="Output field file")
|
|
parser.add_argument("-N","--output_size", type=int, nargs="+", help="Output field size (N0, N1, N2)")
|
|
parser.add_argument("-L","--output_L", type=float, nargs="+", default=None, help="Output field size (L0, L1, L2)")
|
|
parser.add_argument("-dpm","--output_dpm", type=float, nargs="+", default=None, help="Output field dpm (d0, d1, d2)")
|
|
parser.add_argument("-corner","--output_corner", type=float, nargs=3, default=(0.,0.,0.), help="Output field corner (corner0, corner1, corner2)")
|
|
parser.add_argument("-BC","--boundary_conditions", type=int, default=1, help="Boundary conditions (1: periodic, 3: non-periodic)")
|
|
args = parser.parse_args()
|
|
|
|
field_to_field(args.input_file, args.output_file, args.output_size, args.output_L, args.output_dpm, args.output_corner, args.boundary_conditions)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
console_main() |