Python implementation of WSEG-10 fallout model

wseg10.py

import numpy as np
from scipy.special import gamma
from scipy.stats import norm
from affine import Affine

# WSEG-10 fallout model as described in Dan W. Hanifen, "Documentation and Analysis
# of the WSEG-10 Fallout Prediction Model," Thesis, Air Force Institute of Technology,
# March 1980.

class WSEG10:
    def __init__(self, gzx, gzy, yld, ff, wind, wd, shear, tob=0):
        self.translation = ~Affine.translation(gzx, gzy) # translate coordinates relative to GZ (st. mi)
        self.wd = wd # wind direction in degrees (0=N, 90=E, etc.)
        self.yld = yld # yield (MT)
        self.ff = ff # fission fraction, 0 < ff <= 1.0
        self.wind = wind # wind speed (mi/hr)
        self.shear = shear # wind shear in mi/hr-kilofoot
        self.tob = tob # time of burst (hrs)
        # FORTRAN is ugly in any language
        # Store these values in the WSEG10 object to avoid recalculating them
        d = np.log(yld) + 2.42 # physically meaningless, but occurs twice
        self.H_c =  44 + 6.1 * np.log(yld) - 0.205 * abs(d) * d # cloud center height
        lnyield = np.log(yld)
        self.s_0 = np.exp(0.7 + lnyield / 3 - 3.25 / (4.0 + (lnyield + 5.4)**2)) #sigma_0
        self.s_02 = self.s_0**2
        self.s_h = 0.18 * self.H_c # sigma_h
        self.T_c = 1.0573203 * (12 * (self.H_c / 60) - 2.5 * (self.H_c / 60)**2) * (1 - 0.5 * np.exp(-1 * (self.H_c / 25)**2)) # time constant
        self.L_0 = wind * self.T_c # L_0, used by g(x)
        self.L_02 = self.L_0**2
        self.s_x2 = self.s_02 * (self.L_02 + 8 * self.s_02) / (self.L_02 + 2 * self.s_02)
        self.s_x = np.sqrt(self.s_x2) # sigma_x
        self.L_2 = self.L_02 + 2 * self.s_x2
        self.L = np.sqrt(self.L_2) # L
        self.n = (ff * self.L_02 + self.s_x2) / (self.L_02 + 0.5 * self.s_x2) # n
        self.a_1 = 1 / (1 + ((0.001 * self.H_c * wind) / self.s_0)) # alpha_1

    def g(self, x):
        return np.exp(-(np.abs(x) / self.L)**self.n) / (self.L * gamma(1 + 1 / self.n))

    def phi(self, x):
        w = (self.L_0 / self.L) * (x / (self.s_x * self.a_1))
        return norm.cdf(w)

    def D_Hplus1(self, x, y):
        """Returns dose rate at x, y in R/hr at 1 hour after burst. This value includes dose rate from all activity that WILL be deposited at that location, not just that that has arrived by H+1 hr."""
        rx, ry = self.translation * (x, y) * ~Affine.rotation(-self.wd + 270)
        f_x = self.yld * 2e6 * self.phi(rx) * self.g(rx) * self.ff
        s_y = np.sqrt(self.s_02 + ((8 * abs(rx + 2 * self.s_x) * self.s_02) / self.L) + (2 * (self.s_x * self.T_c * self.s_h * self.shear)**2 / self.L_2) + (((rx + 2 * self.s_x) * self.L_0 * self.T_c * self.s_h * self.shear)**2 / self.L**4))
        a_2 = 1 / (1 + ((0.001 * self.H_c * self.wind) / self.s_0) * (1 - norm.cdf(2 * x / self.wind)))
        f_y = np.exp(-0.5 * (ry / (a_2 * s_y))**2) / (2.5066282746310002 * s_y)
        return f_x * f_y
        
    def fallouttoa(self, x):
        """Average time-of-arrival for fallout along hotline at x. Minimum is 0.5hr for any location."""
        T_1 = 1.0
        return np.sqrt(0.25 + (self.L_02 * (x + 2 * self.s_x2) * self.T_c**2) / (self.L_2 * (self.L_02 + 0.5 * self.s_x2)) + ((2 * self.s_x2 * T_1**2) / (self.L_02 + 0.5 * self.s_x)))

    def dose(self, x, y):
        """Estimate of total "Equivalent Residual Dose" (ERD) in R at location x, y from time of fallout arrival to 30 days, including a 90% recovery factor. """
        rx, _ = self.translation * (x, y) * ~Affine.rotation(-self.wd + 270)
        t_a = self.fallouttoa(rx)
        bio = np.exp(-(0.287 + 0.52 * np.log(t_a / 31.6) + 0.04475 * np.log((t_a / 31.6)**2)))
        return self.D_Hplus1(x, y) * bio