diff --git a/run_pm_dmo_NFW_fixed_timestep.md b/run_pm_dmo_NFW_fixed_timestep.md
index 62ea7f0898be1313cc5599ae3089ca6b456bf831..aa39bd2939988f5a7d6f2ed09363e48042a309c7 100644
--- a/run_pm_dmo_NFW_fixed_timestep.md
+++ b/run_pm_dmo_NFW_fixed_timestep.md
@@ -16,86 +16,80 @@ from hotwheels_integrate import *
from hotwheels_io import *
#
-# step 1: config of components
-#
-# at this stage we do not allocate any resource, we just
-# pass the right config parameters to the constructors
-# in order to compile the underlying C libraries.
+# Step 1: Configure components
+# This stage configures components without allocating resources.
+# Configurations are passed to constructors to compile the underlying C libraries.
#
-#call MPI_Init()
-mpi = hwc.MPI().init()
-#configure my malloc with 2GB
-mym = MyMalloc(alloc_bytes=int(2e9))
-#configure the default particle SoA with 1e5 particles
-p = SoA(maxpart=int(1e5), mem=mym)
-#add the particle soa to the multi-type particle SoAs
-soas = SoAs(p, mem=mym)
-# configure the timestep class to go from 0 to 1 Gyr
-# note: G=43007.1 in units of length=kpc, velocity=km/s, mass = 1e10Msun
-# if G=43007.1, masses are in 1e10Msun/h, and radii are in ckpc/h, then this
-# constant converts from Gyr to internal units. Not much to comment about it
+mpi = hwc.MPI().init() # Initialize MPI
+mym = MyMalloc(alloc_bytes=int(2e9)) # Configure memory allocator with 2GB
+p = SoA(maxpart=int(1e5), mem=mym) # Configure P to hold 1e5 particles
+soas = SoAs(p, mem=mym) # Add P to a multi-type SoA container
+# Set up a fixed time-step integrator from 0 to 1 Gyr
+# Conversion factor for Gyr to internal units
gyr_to_cu = 3.086e+16 / (1e9 * 3600 * 24 * 365)
-ts = integrate.FixedTimeStep(soas, G=43007.1, t_from=0., t_to=1. * gyr_to_cu, MPI=mpi)
-ic = NFWIC(rs=100., rho0=1e-6, rs_factor=10.) #init the build of a NFW profile with R<10rs, (note that mass is in units of 1e10Msun)
-# init a refined PM grids with 7 stacked PLACEHIRESREGION at smaller and smaller scales
-# note that PM needs the time-step integrator class TS to attach its DriftTable and kick callbacks
-pm = SuperHiResPM(soas=soas, mem=mym, TS=ts, MPI=mpi, pmgrid=128, grids=8, dt_displacement_factor=0.25)
-#configure to compile all modules in the current folder
-build = make.Build('./', mpi, pm, ts, mym, *soas.values())
-#configure generate SoA headers
-headers = OnTheFly(build.build_name, *build.components, generate_user_c=True)
-
-#
-# step 2: build SoAs headers and compile C files
-#
-
-
-if mpi.rank == 0: #master rank compile and build headers
+ts = integrate.FixedTimeStep(
+ soas,
+ G=43007.1, # Gravitational constant in specific units
+ t_from=0.,
+ t_to=1. * gyr_to_cu,
+ MPI=mpi
+)
+# Initialize a NFW profile with scale radius `rs=100` and density `rho0=1e-6`
+ic = NFWIC(rs=100., rho0=1e-6, rs_factor=10.)
+# Configure a refined PM grid with 7 stacked high-resolution regions
+pm = SuperHiResPM( #wrapper to the PM C library
+ soas=soas,
+ mem=mym,
+ TS=ts, #will use it to attach gravkick callback
+ MPI=mpi,
+ pmgrid=128,
+ grids=8, # number of grids to instantiate
+ dt_displacement_factor=0.25 #factor for DtDisplacement
+)
+build = make.Build('./', mpi, pm, ts, mym, *soas.values()) # Compile all modules in the current directory
+headers = OnTheFly(build.build_name, *build.components, generate_user_c=True) # Generate SoA headers
+
+if mpi.rank == 0: # Master rank handles compilation
headers.write()
build.compile()
-
#
-# step 3: allocate resources (e.g. MPI, MyMalloc, and allocations within it)
+# Step 2: Allocate resources
#
-with (utils.Panic(Build=build) as panic, #attach panic handler to C calls
- utils.Timer(Build=build) as timer, #attach timer handler to C calls
- build.enter(debug=mpi.rank == 0), #parse the compiled object files
- mpi.enter(pm), #attach MPI init info to PM module
- mym.enter(*build.components), #actually allocates the 2GB of ram
- p, #allocate the particle data stracture in the MyMalloc area
- ic.enter(p, mpi.ranks, p.get_maxpart()), #sample the NFW in the particle SoA `p` fields
- pm, #call pm_init() the PM regions
- ts #compute DriftTables if necessary
- ):
-
- pm.compute_accelerations() #first acc computation
-
-
- #
- # step 4: the actual run
- #
-
-
- while ts.time < ts.time_end: #main run.c loop
-
- ts.find_timesteps() #integrator will find timesteps
- #integrator will call halfstep kick, including PM registered drift and kick callbacks
- ts.do_first_halfstep_kick()
- ts.drift() #integretor will drift
- pm.compute_accelerations() #accelerations
- #integrator will call halfstep kick, including PM registered drift and kick callbacks
- ts.do_second_halfstep_kick()
-
- if mpi.rank == 0 and steps % 10 == 0: #sometimes, master rank will do a plot
- f, ax = plt.subplots(1)
- ax.hist2d(p['pos'][:, 0], p['pos'][:, 1], bins=128)
- ax.set_aspect('equal')
- pat = os.path.join(os.getenv('HW_BUILD', '.'), 'snap%d_rank%d.png' % (steps, mpi.rank))
- f.savefig(pat, bbox_inches='tight', dpi=200)
- plt.close(f)
-
-
-print('simualtion finished')
\ No newline at end of file
+with (
+ utils.Panic(Build=build) as panic, # Attach panic handler
+ utils.Timer(Build=build) as timer, # Attach timer handler
+ build.enter(debug=mpi.rank == 0), # Parse compiled objects
+ mpi.enter(pm), # Initialize MPI in the PM module
+ mym.enter(*build.components), # Allocate 2GB memory
+ p, # Allocate particle data structure in MyMalloc
+ ic.enter(p, mpi.ranks, p.get_maxpart()), # Sample NFW profile
+ pm, # Initialize PM and compute first accelerations
+ ts # Compute DriftTables if needed
+):
+
+ #
+ # Step 3: Main simulation loop
+ #
+
+ while ts.time < ts.time_end:
+ ts.find_timesteps() # Determine timesteps
+ ts.do_first_halfstep_kick() # First kick (includes drift/kick callbacks)
+ ts.drift() # Update particle positions
+ pm.compute_accelerations() # Recompute accelerations
+ ts.do_second_halfstep_kick() # Second kick
+
+ # Occasionally, generate plots on the master rank
+ if mpi.rank == 0 and ts.steps % 10 == 0:
+ fig, ax = plt.subplots(1)
+ ax.hist2d(p['pos'][:, 0], p['pos'][:, 1], bins=128)
+ ax.set_aspect('equal')
+ path = os.path.join(
+ os.getenv('HW_BUILD', '.'), f'snap{ts.steps}_rank{mpi.rank}.png'
+ )
+ fig.savefig(path, bbox_inches='tight', dpi=200)
+ plt.close(fig)
+
+print('Simulation finished')
\ No newline at end of file