Resource optimization and monitoring for Parallel Jobs

# File Name - deflection_angle_mpi.py
# This script computes the gravitational deflection angle of light around a massive object
# using MPI for parallelization across multiple processes. Each MPI rank computes a chunk
# of the mass grid, results are gathered at the root process, and a color plot is generated
# to visualize the deflection angles on a logarithmic scale.

# Import MPI module from mpi4py, NumPy for numerical array operations, time for measuring execution time 
# Import os for appending paths, and matplotlib for plotting the results
from mpi4py import MPI
import numpy as np
import time
import os 
import matplotlib.pyplot as plt
import matplotlib.colors as colors

# MPI setup
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

# Constants
G = 6.67430e-11
c = 299792458
M_sun = 1.98847e30

# Parameter grid (same on all ranks)
mass_grid = np.linspace(1, 1000, 10000)  # Solar masses
impact_grid = np.linspace(1e9, 1e12, 10000)  # meters

# Distribute mass grid among ranks
chunk_size = len(mass_grid) // size # Number of elements assigned to each rank
start_idx = rank * chunk_size # Starting index for this rank
end_idx = (rank + 1) * chunk_size if rank != size - 1 else len(mass_grid) # End index (last rank takes the remainder)

local_mass = mass_grid[start_idx:end_idx]
local_result = np.zeros((len(local_mass), len(impact_grid)))

# Timing
local_start = time.time()

# Compute local chunk
for i, M in enumerate(local_mass):
    for j, b in enumerate(impact_grid):
        local_result[i, j] = (4 * G * M * M_sun) / (c**2 * b)

local_end = time.time()
print(f"Rank {rank} local time: {local_end - local_start:.3f} seconds")

# Gather results
result = None
if rank == 0:
    result = np.zeros((len(mass_grid), len(impact_grid)))

comm.Gather(local_result, result, root=0)

if rank == 0:
    total_time = local_end - local_start
    print(f"MPI total time (wall time): {total_time:.3f} seconds")
    result = np.save("result_mpi.npy", result)
    mass_grid = np.save("mass_grid_mpi.npy", mass_grid)
    impact_grid = np.save("impact_grid_mpi.npy", impact_grid)

# Load data
result = np.load("result_mpi.npy")
mass_grid = np.load("mass_grid_mpi.npy")
impact_grid = np.load("impact_grid_mpi.npy")

# Create meshgrid
M, B = np.meshgrid(mass_grid / 1.989e30, impact_grid / 1e9, indexing='ij')

# Create output directory
os.makedirs("plots", exist_ok=True)

plt.figure(figsize=(8,6))
pcm = plt.pcolormesh(B, M, result,
                      norm=colors.LogNorm(vmin=result[result > 0].min(), vmax=result.max()),
                      shading='auto', cmap='plasma')

plt.colorbar(pcm, label='Deflection Angle (radians, log scale)')
plt.xlabel('Impact Parameter (Gm)')
plt.ylabel('Mass (Solar Masses)')
plt.title('Gravitational Deflection Angle - MPI')

plt.tight_layout()
plt.savefig("plots/deflection_angle_mpi.png", dpi=300)
plt.close()

print("MPI plot saved in 'plots/deflection_angle_mpi.png'")

# File Name - deflection_angle_mpi_monitor.py
# This script computes the gravitational deflection angle of light around a massive object
# using MPI for parallelization across multiple processes. Each rank computes a portion
# of the mass grid and the results are gathered at the root process. 
# Additionally, each rank launches a background monitoring thread that periodically logs
# CPU and memory usage during execution. The root process generates a heatmap plot of the
# deflection angle results using a logarithmic color scale.

# Import MPI module from mpi4py, NumPy for numerical array operations, time for measuring execution time
# Import os for appending paths, matplotlib for plotting the results, and psutil for resource monitoring
from mpi4py import MPI
import numpy as np
import time
import os
import psutil
import threading
import matplotlib.pyplot as plt
import matplotlib.colors as colors

# MPI setup
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

# Constants
G = 6.67430e-11
c = 299792458
M_sun = 1.98847e30

# Monitoring setup
pid = os.getpid()
process = psutil.Process(pid)

def monitor(interval=5):
    """Continuously log CPU and memory usage every `interval` seconds."""
    while True:
        cpu = process.cpu_percent(interval=None)   # CPU usage %
        mem = process.memory_info().rss / (1024*1024)  # Memory in MB
        print(f"[Rank {rank}] CPU%: {cpu:.1f} | Mem: {mem:.2f} MB", flush=True)
        time.sleep(interval)

# Start monitoring thread (all ranks or just rank 0)
t = threading.Thread(target=monitor, daemon=True)
t.start()

# Parameter grid (same on all ranks)
mass_grid = np.linspace(1, 1000, 10000)   # Solar masses
impact_grid = np.linspace(1e9, 1e12, 10000)  # meters

# Divide work among ranks
chunk_size = len(mass_grid) // size
start_idx = rank * chunk_size
end_idx = (rank + 1) * chunk_size if rank != size - 1 else len(mass_grid)

local_mass = mass_grid[start_idx:end_idx]
local_result = np.zeros((len(local_mass), len(impact_grid)))

# Timing + Computation
local_start = time.time()

for i, M in enumerate(local_mass):
    for j, b in enumerate(impact_grid):
        local_result[i, j] = (4 * G * M * M_sun) / (c**2 * b)

local_end = time.time()
print(f"[Rank {rank}] Local compute time: {local_end - local_start:.3f} seconds")

# Gather results
result = None
if rank == 0:
    result = np.zeros((len(mass_grid), len(impact_grid)))

comm.Gather(local_result, result, root=0)

# Post-processing + Plotting
if rank == 0:
    total_time = local_end - local_start
    print(f"[Rank 0] MPI total time (wall time): {total_time:.3f} seconds")

    # Save results
    np.save("result_mpi.npy", result)
    np.save("mass_grid_mpi.npy", mass_grid)
    np.save("impact_grid_mpi.npy", impact_grid)

    # Reload for plotting
    result = np.load("result_mpi.npy")
    mass_grid = np.load("mass_grid_mpi.npy")
    impact_grid = np.load("impact_grid_mpi.npy")

    # Create meshgrid
    M, B = np.meshgrid(mass_grid, impact_grid / 1e9, indexing='ij')

    # Create output directory
    os.makedirs("plots", exist_ok=True)

    # Plot results
    plt.figure(figsize=(8,6))
    pcm = plt.pcolormesh(B, M, result,
                         norm=colors.LogNorm(vmin=result[result > 0].min(), vmax=result.max()),
                         shading='auto', cmap='plasma')

    plt.colorbar(pcm, label='Deflection Angle (radians, log scale)')
    plt.xlabel('Impact Parameter (Gm)')
    plt.ylabel('Mass (Solar Masses)')
    plt.title('Gravitational Deflection Angle - MPI with Monitoring')

    plt.tight_layout()
    plt.savefig("plots/deflection_angle_mpi.png", dpi=300)
    plt.close()

    print("MPI plot saved in 'plots/deflection_angle_mpi.png'")