main.py

# This module runs the code for assignment 1 of Stochastic Simulation

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

import time
from tqdm import trange

# --------------
# Volume dimensions
# --------------
def sphere(x, y, z, k):
    """Checks if the point is within the sphere and passes True if so."""
    if k <= 0:
        raise ValueError(f"k needs to be > 0. You got k: {k}")

    # Sphere dimensions
    hits = x*x + y*y + z*z <= k ** 2

    return hits


def torus(x, y, z, R, r):
    """Checks if the point is within the torus and passes True if so."""
    if R <= 0 or r <= 0:
        raise ValueError(f"R and r need to be > 0." 
                         f"You got R: {R} and r: {r}")

    # Torus dimensions
    hits = (np.sqrt(x*x + y*y) - R) ** 2 + z*z <= r ** 2
    
    return hits

# --------------
# Exact solution
# --------------

def volumeAnalytical(k, R, r):
    # only works for centered toroid

    R_1 = (R ** 2 + k ** 2 - r ** 2) / (2 * R)

    F_1 = lambda x, a : -np.power(a ** 2 - x ** 2, 1.5) / 3
    F_2 = lambda x, a : (np.arcsin(x / a) + 0.5 * np.sin(2 * np.arcsin(x / a))) * 0.5 * a ** 2

    V = F_1(k, k) - F_1(R_1, k)
    V += F_1(R_1 - R, r) - F_1(-r, r)
    V += R * (F_2(R_1 - R, r) - F_2(-r, r))

    return V * 4 * np.pi

# --------------
# Sampling
# --------------
def uniformrandom(N, seed=None):
    np.random.seed(seed=seed)
    x = np.random.rand(N)
    y = np.random.rand(N)
    z = np.random.rand(N)

    return np.array([x, y, z])


def deterministic_XYZ(N, seed=None):
    sequence = np.empty((3, N))
    m = 3.8

    # define a self-contained region
    b = m * 0.5 * (1 - 0.5)
    a = m * b * (1 - b)
    
    # 'seed' the deterministic sequence
    np.random.seed(seed=seed)
    sequence[:, 0] = np.random.uniform(a, b, 3)
    
    for i in range(1, N):
        sequence[:, i] = m * sequence[:, i - 1] * (1 - sequence[:, i - 1])

    # normalise [a, b] -> [0, 1]
    sequence = (sequence - a) / (b - a)

    return sequence 


# define random number generator

# --------------
# Errors
# --------------
def standard_error(sample_std, N):
    """Calculates the standard error of the mean."""
    return sample_std / np.sqrt(N)


# --------------
# monte carlo
# --------------
def montecarlo(prng, k=1, R=0.75, r=0.4, throws=100000, 
               xc=0, yc=0, zc=0, 
               plot=False,
               mode='uniform',
               p_b=0.5,
               b_radius=1.1,
               s_radius=0.4):   # 'uniform' for Q1&Q3a, 'mixture' for Q3b

    if mode == 'uniform':
        rand = prng(throws)
        # normalise [0, 1] -> [-radius, radius]
        x, y, z = b_radius * (np.ones((3, throws)) - 2 * rand)

    elif mode == 'mixture':
        randnum_b = prng(throws)
        randnum_s = prng(throws)
        
        from_b = np.random.rand(throws) < p_b # boolean array, True = from B box, False = from S box
        
        s_center = np.array([xc, yc, zc])

        # empty arrays to hold the random points from each box
        x = np.empty(throws)
        y = np.empty(throws)
        z = np.empty(throws)

        # normalise [0, 1] -> [-b_radius, b_radius]
        x[from_b] = b_radius * (1 - 2 * randnum_b[0, from_b])
        y[from_b] = b_radius * (1 - 2 * randnum_b[1, from_b])
        z[from_b] = b_radius * (1 - 2 * randnum_b[2, from_b])

        # normalise [0, 1] -> [-s_radius, s_radius]
        x[~from_b] = b_radius * (1 - 2 * randnum_s[0, ~from_b])
        y[~from_b] = b_radius * (1 - 2 * randnum_s[1, ~from_b])
        z[~from_b] = b_radius * (1 - 2 * randnum_s[2, ~from_b])


    # check hits
    sphereHits = sphere(x, y, z, k)
    torusHits = torus(x-xc, y-yc, z-zc, R, r)

    totalHits = np.sum(sphereHits & torusHits)

    # calculate intersection volume
    box_volume = (2 * b_radius) ** 3
    intersection_volume = box_volume * (totalHits / throws)

    # plotting
    if plot == True:
        plotintersection(x, y, z, sphereHits, torusHits, b_radius)

    return intersection_volume, totalHits


def run_monte_carlo(
        N=100, 
        prng=uniformrandom, 
        k=1, 
        R=0.75, 
        r=0.4, 
        throws=100000,
        xc=0, yc=0,zc=0,
        mode='uniform',
        p_b=0.5,
        b_radius=1.1,
        s_radius=0.4):
    """Runs the monte carlo simulation N times."""

    all_volumes = []
    all_hits = []

    # Time progress bar
    t0 = time.perf_counter()
    for _ in trange(N, desc="Monte Carlo runs", leave=False):
        intersection_volume, hits = montecarlo(
            prng, k, R, r, throws, 
            xc, yc, zc,
            plot=False,
            mode=mode,
            p_b=p_b,
            b_radius=b_radius,
            s_radius=s_radius
        )
        all_volumes.append(intersection_volume)
        all_hits.append(hits)

    t1 = time.perf_counter()

    sample_std = np.std(all_volumes)
    average_volume = np.mean(all_volumes)

    elapsed = t1 - t0
    print(f"for b_radius={b_radius}, k={k}, R={R}, r={r}, and throws={throws}, we get:")
    print(f"average_volume: {average_volume}, sample variance: {sample_std}")
    print(f"Elapsed: {elapsed:.3f}s  ({elapsed/N:.6f}s per run)")

    return sample_std, average_volume

# --------------
# Plotting code
# --------------
def plotintersection(x, y, z, sphereHits, torusHits, radius):
    # store points for each category
    
    # Add code
    intersection = sphereHits & torusHits

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # plot intersection
    ax.scatter(
        x[intersection],
        y[intersection],
        z[intersection],
        s=6, alpha=0.7, label=f"Intersection ({intersection.sum()})"
    )

    # labels and title
    ax.set_xlabel('X Axis')
    ax.set_ylabel('Y Axis')
    ax.set_zlabel('Z Axis')
    ax.set_title("Intersection")
    ax.legend()
    
    # set limits for each axis
    ax.set_xlim([-radius, radius])
    ax.set_ylim([-radius, radius])
    ax.set_zlim([-radius, radius])

    plt.show()

# plot histogram for deterministic sequence -> show that it is not uniform
def plotDeterministicHistogram(N):
    xDet, _, _ = deterministic_XYZ(N)
    xRand, _, _ = uniformrandom(N)

    plt.hist(xDet, density=True, label="deterministic")
    plt.hist(xRand, density=True, histtype="step", label="numpy prng")
    plt.xlabel("Sample value")
    plt.ylabel("Density of occurence")
    plt.legend()
    plt.show()
    
    return
    

# Plot error changes and estimates
def convergencePlot(N, b_radius, k, R, r, maxThrows, throwsSamples):
    throwsList = np.logspace(1, np.log(maxThrows) / np.log(10), throwsSamples, base=10, dtype=np.int32)

    means = np.empty((2, throwsSamples))
    stds = np.empty((2, throwsSamples))
    for i in trange(throwsSamples, desc="Convergence plot: ", leave=False):
        throws = throwsList[i]
        vols = np.empty((2, N))
        for n in range(N):
            vols[0, n], _ = montecarlo(uniformrandom, b_radius, k, R, r, throws)
            vols[1, n], _ = montecarlo(deterministic_XYZ, b_radius, k, R, r, throws)

        means[:, i] = np.mean(vols, axis=1)
        stds[:, i] = np.std(vols, axis=1)

    fig, ax = plt.subplots()

    exactVolume = volumeAnalytical(k, R, r)

    ax.errorbar(throwsList, means[0, :], yerr=stds[0, :], fmt='o', label="uniform")
    ax.errorbar(throwsList, means[1, :], yerr=stds[1, :], fmt='o', label="deterministic")
    ax.hlines(exactVolume, np.min(throwsList), np.max(throwsList), linestyle="dashed", color="grey", label="exact")
    ax.set_xscale("log")
    ax.set_xlabel("Amount of throws")
    ax.set_ylabel("Volume estimate (a.u.)")
    ax.legend()

    plt.show()

    return

def line_plots(results_volume, results_error, p_values, s_radii):
    plt.figure(figsize=(14, 6))

    # plot 1: biased estimated volume
    plt.subplot(1, 2, 1)
    
    for i, s_rad in enumerate(s_radii):
        plt.plot(p_values, results_volume[i, :], 'o-', label=f's_radius = {s_rad:.2f}')
    
    plt.xlabel("probability from box B)")
    plt.ylabel("estimated Volume (biased)")
    plt.title("volume vs. p_b (grouped by small box radius)")
    plt.legend()
    plt.grid(True)

    # plot 2: std dev error
    # We can use this to find the best (p_b, s_radius) combination.
    plt.subplot(1, 2, 2)
    
    for i, s_rad in enumerate(s_radii):
        plt.plot(p_values, results_error[i, :], 'o-', label=f's_radius = {s_rad:.2f}')
    
    plt.xlabel("probability from box b")
    plt.ylabel("sample standard deviation (error)")
    plt.title("error vs. p_b (grouped by small box radius)")
    plt.legend()
    plt.grid(True)
    
    plt.tight_layout()
    plt.show()

# --------------
# Running code
# --------------

def main():
    # case a:

    a_sample_std, a_average_volume = run_monte_carlo(
        N=100, 
        prng=uniformrandom, 
        b_radius=1.1, 
        k=1, 
        R=0.75, 
        r=0.4, 
        throws=int(1e5))

    # case b:
    b_sample_std, b_average_volume = run_monte_carlo(
        N=100, 
        prng=uniformrandom, 
        b_radius=1.1, 
        k=1, 
        R=0.5, 
        r=0.5, 
        throws=int(1e5))
 
    # plot case a
    montecarlo(prng=uniformrandom, 
               b_radius=1.1, 
               k=1, 
               R=0.75, 
               r=0.4, 
               throws=int(1e4), 
               plot=True)

    # convergence plot case a
    convergencePlot(N=100,
                    b_radius=1.1,
                    k=1,
                    R=0.75,
                    r=0.4,
                    maxThrows=int(1e5),
                    throwsSamples=20)
    
    # convergence plot case b
    convergencePlot(N=100,
                    b_radius=1.1,
                    k=1,
                    R=0.5,
                    r=0.5,
                    maxThrows=int(1e5),
                    throwsSamples=20)

    plotDeterministicHistogram(int(1e5))

def test_q3b():
    n_throws = int(1e5)
    b_radius = 1.1         
    
    k_3, R_3, r_3 = 1.0, 0.75, 0.4
    xc_3, yc_3, zc_3 = 0, 0, 0.1
 
    p_values = np.array([0.1, 0.3, 0.5, 0.7, 0.9])
    s_radii = np.array([0.1, 0.3, 0.5, 0.7, 0.9])
    
    n_repeats = 500

    results_avg_volume = np.empty((len(s_radii), len(p_values)))
    results_std_dev = np.empty((len(s_radii), len(p_values)))
    
    for i, s_rad in enumerate(s_radii):
        for j, p in enumerate(p_values):
            
            print(f"    > testing s_rad={s_rad:.2f}, p_b={p:.1f} ...")

            std_dev, avg_vol = run_monte_carlo(
                N=n_repeats,
                prng=uniformrandom,
                k=k_3, R=R_3, r=r_3,
                throws=n_throws,
                xc=xc_3, yc=yc_3, zc=zc_3,
                mode='mixture',
                p_b=p,
                b_radius=b_radius,
                s_radius=s_rad
            )
            
            # store results
            results_avg_volume[i, j] = avg_vol
            results_std_dev[i, j] = std_dev
    
    line_plots(results_avg_volume, results_std_dev, p_values, s_radii)

main()
# test_q3b()