import numpy as np import g3read as g3 from shutil import copyfile import matplotlib.pyplot as plt import matplotlib as mpl mpl.use('Agg') # batch / no GUI import os def setup_nozzle_3d( fin="glass_10x10x10", fout="box.ic", Ndup=40, Ndup2=2, gamma=5.0/3.0, rho=0.125, P=1.0, ic_header_path="ic_header", preview_png="nozzle_3d_geometry.png", ): # ------------------------------------------------------------------ # Read glass and normalize to unit cube # ------------------------------------------------------------------ f_in = g3.GadgetFile(fin) xin = f_in.read_new("POS ", 0).astype(np.float32) # header / boxsize try: h_in = f_in.read_new("HEAD", 0) except Exception: h_in = getattr(f_in, "header", {}) if isinstance(h_in, dict): L_in = float(h_in.get("BOXSIZE", 1.0)) else: # try common names if hasattr(h_in, "boxsize"): L_in = float(h_in.boxsize) elif hasattr(h_in, "BoxSize"): L_in = float(h_in.BoxSize) else: L_in = 1.0 xin /= L_in Nin = int(np.sum(getattr(f_in, "header").npart)) # ------------------------------------------------------------------ # Tiling parameters # ------------------------------------------------------------------ L = Ndup / Ndup2 # physical length in x print("Length of the region to fill is", L) Ntiles = Ndup * Ndup2 * Ndup2 Np = Ntiles * Nin # ------------------------------------------------------------------ # Build positions by tiling: (xin + (ix,iy,iz)) / Ndup2 # ------------------------------------------------------------------ ix, iy, iz = np.meshgrid( np.arange(Ndup, dtype=np.float32), np.arange(Ndup2, dtype=np.float32), np.arange(Ndup2, dtype=np.float32), indexing="ij" ) offsets = np.stack([ix, iy, iz], axis=-1).reshape(-1, 3) # (Ntiles,3) # Broadcast and reshape -> (Np,3) x = ((xin[None, :, :] + offsets[:, None, :]) / float(Ndup2)).reshape(-1, 3).astype(np.float32) # ------------------------------------------------------------------ # Fields: velocity, internal energy, mass, ID # ------------------------------------------------------------------ v = np.zeros((Np, 3), dtype=np.float32) u = np.zeros(Np, dtype=np.float32) m = np.zeros(Np, dtype=np.float32) # IDL: id = lindgen(Np) → 0..Np-1 ids = np.arange(Np, dtype=np.int32) u[:] = P / (gamma - 1.0) / rho # mpart = (rho * Ndup2^3 * BoxSize^3) / Np, with BoxSize=1 mpart = (rho * (Ndup2**3) * 1.0**3) / Np m[:] = mpart # ------------------------------------------------------------------ # Mark walls and nozzle region # ------------------------------------------------------------------ # 1) Wall at end of slab: x in (L - 0.5, L) dL = 0.5 mask_end = x[:, 0] > (L - dL) ids[mask_end] = 0 # 2) Nozzle in left part: length dL = 8 x0 = 0.5 dL_noz = 8.0 dmin = 0.15 amp = (0.5 - dmin) / 2.0 mask_noz_region = (x[:, 0] > x0) & (x[:, 0] < x0 + dL_noz) if np.any(mask_noz_region): dy = np.abs(x[:, 1] - 0.5) theta = (x[:, 0] - x0) / dL_noz * 2.0 * np.pi + np.pi / 2.0 dy_nozzle = 0.5 + amp * (np.sin(theta) - 1.0) mask_outside = mask_noz_region & (dy > dy_nozzle) ids[mask_outside] = 0 # 3) Wall at beginning of slab: x < 0.5, id=0, give vx=-1 dL = 0.5 dv = -1.0 mask_start = x[:, 0] < dL ids[mask_start] = 0 v[mask_start, 0] = dv # 4) Shift everything to the right by L x[:, 0] += L # ------------------------------------------------------------------ # Write Gadget IC using a reference header file # ------------------------------------------------------------------ if not os.path.exists(ic_header_path): raise FileNotFoundError(f"Reference header '{ic_header_path}' not found.") copyfile(ic_header_path, fout) f_out = g3.GadgetFile(fout, is_snap=False) # Header fields as in previous scripts / IDL logic f_out.header.npart = [Np, 0, 0, 0, 0, 0] f_out.header.npartTotal = [Np, 0, 0, 0, 0, 0] f_out.header.time = 0.0 f_out.header.boxsize = 1.0 try: f_out.header.massarr = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0] except Exception: pass f_out.write_header(f_out.header, filename=fout) # Add file blocks f_out.add_file_block('POS ', Np * 4 * 3, partlen=4 * 3) f_out.add_file_block('VEL ', Np * 4 * 3, partlen=4 * 3) f_out.add_file_block('ID ', Np * 4, partlen=4, dtype=np.int32) f_out.add_file_block('MASS', Np * 4, partlen=4) f_out.add_file_block('U ', Np * 4, partlen=4) # Write data f_out.write_block('POS ', 0, x, filename=fout) f_out.write_block('VEL ', 0, v, filename=fout) f_out.write_block('ID ', 0, ids, filename=fout) f_out.write_block('MASS', 0, m, filename=fout) f_out.write_block('U ', 0, u, filename=fout) print(f"Wrote '{fout}' with Np={Np}, L={L}") # ------------------------------------------------------------------ # Quick geometry plot, x–y projection # ------------------------------------------------------------------ try: fig, ax = plt.subplots(figsize=(10, 3)) # all particles ax.plot(x[:, 0], x[:, 1], 'o', ms=0.5, color='0.6') # fixed / boundary particles (id == 0) mask_fixed_any = (ids == 0) if np.any(mask_fixed_any): ax.plot(x[mask_fixed_any, 0], x[mask_fixed_any, 1], 'o', ms=3, color='C0') # moving wall particles (id == 0 and v_x != 0) mask_moving = (ids == 0) & (v[:, 0] != 0.0) if np.any(mask_moving): ax.plot(x[mask_moving, 0], x[mask_moving, 1], 'o', ms=3, marker='x', color='C3') ax.set_xlim(0, 40) ax.set_xlabel("x") ax.set_ylabel("y") ax.set_aspect('equal', adjustable='box') fig.tight_layout() fig.savefig(preview_png, dpi=150, bbox_inches='tight') plt.close(fig) print(f"Saved geometry preview to '{preview_png}'") except Exception as e: print(f"[preview skipped] {e}") if __name__ == "__main__": setup_nozzle_3d()