Diffusion approximation

diffusion.py

import math
import numpy as np
import matplotlib.pyplot as plt

class Color(object):
    def __init__(self, r = 0.0, g = 0.0, b = 0.0):
        if not isinstance(r, float) or \
           not isinstance(g, float) or \
           not isinstance(b, float):
            raise Exception('Bad arguments!!')
        self.r = r
        self.g = g
        self.b = b

    def __add__(self, c):
        return Color(self.r + c.r, self.g + c.g, self.b + c.b)

    def __neg__(self):
        return Color(-self.r, -self.g, -self.b)

    def __sub__(self, c):
        return self.__add__(c.__neg__())

    def __mul__(self, c):
        if isinstance(c, float):
            return Color(self.r * c, self.g * c, self.b * c)
        else:
            return Color(self.r * c.r, self.g * c.g, self.b * c.b)

    def __rmul__(self, c):
        if isinstance(c, float):
            return Color(self.r * c, self.g * c, self.b * c)
        else:
            return Color(self.r * c.r, self.g * c.g, self.b * c.b)

    def __truediv__(self, c):
        if isinstance(c, float):
            return Color(self.r / c, self.g / c, self.b / c)
        else:
            return Color(self.r / c.r, self.g / c.g, self.b / c.b)

    def __str__(self):
        return '[ %f, %f, %f ]' % (self.r, self.g, self.b)

    @staticmethod
    def zeros():
        return Color(0.0, 0.0, 0.0)

    @staticmethod
    def ones():
        return Color(1.0, 1.0, 1.0)

    @staticmethod
    def sqrt(c):
        return Color(math.sqrt(c.r), math.sqrt(c.g), math.sqrt(c.b))

    @staticmethod
    def exp(c):
        return Color(math.exp(c.r), math.exp(c.g), math.exp(c.b))

    @staticmethod
    def erf(c):
        return Color(math.erf(c.r), math.erf(c.g), math.erf(c.b))

class Dipole(object):
    def __init__(self, sigmap_s, sigma_a, eta):
        A = (1.0 + self.Fdr(eta)) / (1.0 - self.Fdr(eta))

        self.sigmap_t = sigma_a + sigmap_s
        self.sigma_tr = Color.sqrt(3.0 * sigma_a * self.sigmap_t)
        self.alphap   = sigmap_s / self.sigmap_t
        self.zr       = Color(1.0, 1.0, 1.0) / self.sigmap_t
        self.zv       = -self.zr * (1.0 + (4.0 / 3.0) * A)

    def Rd(self, r):
        r2 = r * r
        dr = Color.sqrt(Color(r2, r2, r2) + self.zr * self.zr)
        dv = Color.sqrt(Color(r2, r2, r2) + self.zv * self.zv)
        ret = (self.alphap / (4.0 * math.pi)) * \
              ((self.zr * (dr * self.sigma_tr + Color(1.0, 1.0, 1.0)) * \
                Color.exp(-self.sigma_tr * dr)) / (dr * dr * dr) - \
               (self.zv * (dv * self.sigma_tr + Color(1.0, 1.0, 1.0)) * \
                Color.exp(-self.sigma_tr * dv)) / (dv * dv * dv))
        return ret

    def Fdr(self, eta):
        if eta >= 1.0:
            return -1.4399 / (eta * eta) + 0.7099 / eta + 0.6681 + \
                    0.0636 * eta
        else:
            return -0.4399 + 0.7099 / eta - 0.3319 / (eta * eta) + \
                    0.0636 / (eta * eta * eta)

class Multipole(object):
    def __init__(self, sigmap_s, sigma_a, eta, d=0.1, n=2):
        A = (1.0 + self.Fdr(eta)) / (1.0 - self.Fdr(eta))

        self.sigmap_t = sigma_a + sigmap_s
        self.sigma_tr = Color.sqrt(3.0 * sigma_a * self.sigmap_t)
        self.alphap   = sigmap_s / self.sigmap_t
        self.l        = Color(1.0, 1.0, 1.0) / self.sigmap_t
        self.zb       = self.l * (2.0 / 3.0) * A

        self.d = d
        self.n = n

    def Rd(self, r):
        ret = Color(0.0, 0.0, 0.0)
        for i in range(-self.n, self.n+1):
            zr = 2.0 * i * (Color(self.d, self.d, self.d) + 2.0 * self.zb) + self.l
            zv = 2.0 * i * (Color(self.d, self.d, self.d) + 2.0 * self.zb) - self.l - 2.0 * self.zb
            r2 = r * r
            dr = Color.sqrt(Color(r2, r2, r2) + zr * zr)
            dv = Color.sqrt(Color(r2, r2, r2) + zv * zv)
            ret += (self.alphap / (4.0 * math.pi)) * \
                   ((zr * (dr * self.sigma_tr + Color(1.0, 1.0, 1.0)) * \
                     Color.exp(-self.sigma_tr * dr)) / (dr * dr * dr) - \
                    (zv * (dv * self.sigma_tr + Color(1.0, 1.0, 1.0)) * \
                     Color.exp(-self.sigma_tr * dv)) / (dv * dv * dv))
        return ret

    def Fdr(self, eta):
        if eta >= 1.0:
            return -1.4399 / (eta * eta) + 0.7099 / eta + 0.6681 + \
                    0.0636 * eta
        else:
            return -0.4399 + 0.7099 / eta - 0.3319 / (eta * eta) + \
                    0.0636 / (eta * eta * eta)

class QuantizedDiffusion(object):
    def __init__(self, sigmap_s, sigma_a, eta, d=0.1, n=2, t1=0.1, k=20):
        self.sigmap_s = sigmap_s
        self.sigma_a = sigma_a
        self.sigmap_t = sigma_a + sigmap_s
        self.alphap = self.sigmap_s / self.sigmap_t
        self.D = (2.0 * self.sigma_a + self.sigmap_s) / (3.0 * self.sigmap_t * self.sigmap_t)
        self.eta = eta
        self.zb = 2.0 * self.A(self.eta) * self.D
        self.d = d
        self.n = n
        self.k = k

        s = 1.618
        self.tau = [ t1 * (s ** (i - 1)) for i in range(self.k + 1) ]

    def w(self, i):
        return (Color.exp(-self.tau[i] * self.sigma_a) - Color.exp(-self.tau[i+1] * self.sigma_a)) / self.sigma_a

    def G2D(self, v, r):
        rv = Color.ones() * r
        return (Color.ones() / (2.0 * math.pi * v)) * Color.exp(- rv * rv / (2.0 * v))

    def A(self, eta):
        return (1.0 + 3.0 * self.C2(self.eta)) / (1.0 - 2.0 * self.C1(self.eta))

    def Rd(self, r):
        ret = Color.zeros()
        for i in range(self.k):
            vi = self.D * (self.tau[i] + self.tau[i+1])
            ret += self.w_R(vi) * self.w(i) * self.G2D(vi, r)

        return self.alphap * ret

    def w_R(self, vi):
        ret = Color.zeros()
        zv = Color.zeros()
        dv = Color.ones() * self.d
        for j in range(-self.n, self.n+1):
            m_r = 2.0 * j * (dv + 2.0 * self.zb)
            m_v = 2.0 * j * (dv + 2.0 * self.zb) - 2.0 * self.zb
            w_phiR = self.w_phi(vi, zv, dv, m_r) + self.w_phi(vi, zv, dv, -m_v)
            w_ER   = self.w_E(vi, zv, dv, m_r)   - self.w_E(vi, zv, dv, -m_v)
            ret += self.C_phi(self.eta) * w_phiR + self.C_E(self.eta) * w_ER

        return ret

    def w_phi(self, v, z1, z2, m):
        a1 = (m + self.sigmap_t * v + z1) / Color.sqrt(2.0 * v)
        a2 = (m + self.sigmap_t * v + z2) / Color.sqrt(2.0 * v)
        return (self.alphap * self.sigmap_t * 0.5) * \
               Color.exp(m * self.sigmap_t + 0.5 * self.sigmap_t * self.sigmap_t * v) * \
               (Color.erf(a2) - Color.erf(a1))

    def w_E(self, v, z1, z2, m):
        a1 = (m * m) / (2.0 * v) + (m + self.sigmap_t * v) * z1 / v + (z1 * z1) / (2.0 * v)
        a2 = (m * m) / (2.0 * v) + (m + self.sigmap_t * v) * z2 / v + (z2 * z2) / (2.0 * v)
        return self.D * self.sigmap_t * \
               (-self.w_phi(v, z1, z2, m) + self.alphap * (Color.exp(-a1) - Color.exp(-a2)) / Color.sqrt(2.0 * math.pi * v))

    def C_phi(self, eta):
        return 0.25 * (1.0 - 2.0 * self.C1(eta))

    def C_E(self, eta):
        return 0.5 * (1.0 - 3.0 * self.C2(eta))

    def C1(self, eta):
        eta2 = eta * eta
        eta3 = eta2 * eta
        eta4 = eta3 * eta
        eta5 = eta4 * eta
        if eta < 1.0:
            return (0.919317 - 3.4793 * eta + 6.75335 * eta2 \
                    -7.80989 * eta3 + 4.98554 * eta4 - 1.36881 * eta5) / 2.0
        else:
            return (-9.23372 + 22.2272 * eta - 20.9292 * eta2 + 10.2291 * eta3 \
                    - 2.54396 * eta4 + 0.254913 * eta5) / 2.0

    def C2(self, eta):
        eta2 = eta * eta
        eta3 = eta2 * eta
        eta4 = eta3 * eta
        eta5 = eta4 * eta
        if eta < 1.0:
            return (0.828421 - 2.62051 * eta + 3.36231 * eta2 - 1.95284 * eta3 \
                    + 0.236494 * eta4 + 0.145787 * eta5) / 3.0
        else:
            return (-1641.1 + 135.926 / eta3 - 656.175 / eta2 + 1376.53 / eta \
                    + 1213.67 * eta - 568.556 * eta2 + 164.798 * eta3 \
                    - 27.0181 * eta4 + 1.91826 * eta5) / 3.0

class BetterDipole(object):
    def __init__(self, sigmap_s, sigma_a, eta):
        self.sigmap_t = sigmap_s + sigma_a
        self.alphap   = sigmap_s / self.sigmap_t
        self.eta = eta
        self.A = (1.0 + 3.0 * self.C2(eta)) / (1.0 - 2.0 * self.C1(eta))
        self.D = (2.0 * sigma_a + sigmap_s) / (3.0 * self.sigmap_t * self.sigmap_t)
        self.sigma_tr = Color.sqrt(sigma_a / self.D)
        self.zr = Color.ones() / self.sigmap_t
        self.zb = 2.0 * self.A * self.D
        self.zv = -self.zr - 2.0 * self.zb
        self.Cphi = Color.ones() * (0.25 * (1.0 - 2.0 * self.C1(eta)))
        self.CE   = Color.ones() * (0.50 * (1.0 - 3.0 * self.C2(eta)))

    def Rd(self, r):
        r2 = r * r
        dr = Color.sqrt(Color.ones() * r2 + self.zr * self.zr)
        dv = Color.sqrt(Color.ones() * r2 + self.zv * self.zv)
        ar = self.zr * (self.sigma_tr * dr + Color.ones()) / (dr * dr)
        av = self.zv * (self.sigma_tr * dv + Color.ones()) / (dv * dv)
        br = Color.exp(-self.sigma_tr * dr) / dr
        bv = Color.exp(-self.sigma_tr * dv) / dv
        return (self.alphap * self.alphap / (4.0 * math.pi)) * \
               ((self.CE * ar + self.Cphi / self.D) * br - \
                (self.CE * av + self.Cphi / self.D) * bv)

    def C1(self, eta):
        eta2 = eta * eta
        eta3 = eta2 * eta
        eta4 = eta3 * eta
        eta5 = eta4 * eta
        if eta < 1.0:
            return (0.919317 - 3.4793 * eta + 6.75335 * eta2 \
                    -7.80989 * eta3 + 4.98554 * eta4 - 1.36881 * eta5) / 2.0
        else:
            return (-9.23372 + 22.2272 * eta - 20.9292 * eta2 + 10.2291 * eta3 \
                    - 2.54396 * eta4 + 0.254913 * eta5) / 2.0

    def C2(self, eta):
        eta2 = eta * eta
        eta3 = eta2 * eta
        eta4 = eta3 * eta
        eta5 = eta4 * eta
        if eta < 1.0:
            return (0.828421 - 2.62051 * eta + 3.36231 * eta2 - 1.95284 * eta3 \
                    + 0.236494 * eta4 + 0.145787 * eta5) / 3.0
        else:
            return (-1641.1 + 135.926 / eta3 - 656.175 / eta2 + 1376.53 / eta \
                    + 1213.67 * eta - 568.556 * eta2 + 164.798 * eta3 \
                    - 27.0181 * eta4 + 1.91826 * eta5) / 3.0

def dipole_plot(sigmap_s, sigma_a, eta):
    # Prepare dipole
    dipole = Dipole(sigmap_s, sigma_a, eta)

    # Plot
    xs = [ x * 0.01 for x in range(500) ]
    ys = [ Color.sqrt(2.0 * math.pi * x * dipole.Rd(x)) for x in xs ]

    rs = [ y.r for y in ys ]
    gs = [ y.g for y in ys ]
    bs = [ y.b for y in ys ]

    return xs, (rs, gs, bs)

def multipole_plot(sigmap_s, sigma_a, eta):
    # Prepare dipole
    multipole = Multipole(sigmap_s, sigma_a, eta, d=10.0, n=2)

    # Plot
    xs = [ x * 0.01 for x in range(500) ]
    ys = [ Color.sqrt(2.0 * math.pi * x * multipole.Rd(x)) for x in xs ]

    rs = [ y.r for y in ys ]
    gs = [ y.g for y in ys ]
    bs = [ y.b for y in ys ]

    return xs, (rs, gs, bs)

def gamma_plot():
    sigmap_s = Color(1.00, 1.00, 1.00)
    sigma_a  = Color(0.005, 0.50, 0.50)
    eta      = 1.5

    dp = Dipole(sigmap_s, sigma_a, eta)
    qd = QuantizedDiffusion(sigmap_s, sigma_a, eta, d=1.0e8, n=0, t1=1.0e-4, k=30)
    bd = BetterDipole(sigmap_s, sigma_a, eta)

    xs = [ x * 0.1 for x in range(500) ]
    yd = [ math.sqrt(2.0 * math.pi * x * dp.Rd(x).r) for x in xs ]
    yq = [ math.sqrt(2.0 * math.pi * x * qd.Rd(x).r) for x in xs ]
    yb = [ math.sqrt(2.0 * math.pi * x * bd.Rd(x).r) for x in xs ]

    fig, ax = plt.subplots()
    l1, = ax.plot(xs, yd, 'r', label='Dipole')
    l2, = ax.plot(xs, yq, 'g', label='QD')
    l3, = ax.plot(xs, yb, 'b', label='Better DP')
    ax.set_title('Compare models', fontsize=20)
    ax.set_xlabel('Distance [mm]', fontsize=20)
    ax.set_ylabel('Reflectance $2 \pi R_d(r)$', fontsize=20)
    ax.legend(handles=[l1, l2, l3], loc=1)
    plt.show()

def simple_plot():
    # Jensen's Skin2
    sigmap_s = Color(1.09, 1.59, 1.79)
    sigma_a  = Color(0.013, 0.070, 0.145)
    eta      = 1.3

    xs, (rs, gs, bs) = dipole_plot(sigmap_s, sigma_a, eta)
    xm, (rm, gm, bm) = multipole_plot(sigmap_s, sigma_a, eta)

    fig, ax = plt.subplots()
    ax.plot(xs, rs, 'r-.')
    ax.plot(xs, gs, 'g-.')
    ax.plot(xs, bs, 'b-.')
    ax.plot(xm, rm, 'r--')
    ax.plot(xm, gm, 'g--')
    ax.plot(xm, bm, 'b--')
    plt.show()

def main():
    gamma_plot()

if __name__ == '__main__':
    main()