richards_objfun.py

import sys
import numpy as np
import scipy.sparse as sp
from SimPEG import Mesh, Maps, DataMisfit
from SimPEG.FLOW import Richards

plotIt = False

order = ['Ks', 'theta_r', 'theta_s', 'alpha', 'n']
soil = Richards.Empirical.VanGenuchtenParams().sandy_clay_loam
model = np.r_[[soil[o] for o in order]]

n = 40
m = 40

couples = [
    [0, 1], [0, 2],
    [0, 3], [0, 4],
    [1, 2], [1, 3],
    [1, 4], [2, 3],
    [2, 4], [3, 4]
]

ksrange = np.logspace(-7, -5, n)
trrange = np.linspace(0.01, 0.1, m)
tsrange = np.linspace(0.25, 0.45, m)
arange = np.linspace(2.5, 4.5, m)
nrange = np.linspace(1.1, 1.4, m)


def get_problem(rx_type):
    M = Mesh.TensorMesh([np.ones(40)/100], x0='N')
    M.setCellGradBC('dirichlet')
    bc = np.array([-41.5, -5])/100
    h = np.zeros(M.nC) + bc[0]

    layer1 = np.ones(40)
    P = sp.csr_matrix(np.kron(np.eye(5), np.c_[layer1, ]))
    pmap = Maps.Projection(P.shape[1], P.indices)
    k_fun, theta_fun = Richards.Empirical.van_genuchten(M)
    # ks_1
    wires = Maps.Wires(*zip(order, [40]*len(order)))

    # Maps.ExpMap(M) * if you want log space
    k_fun.KsMap = wires.Ks * pmap
    k_fun.alphaMap = wires.alpha * pmap
    theta_fun.alphaMap = wires.alpha * pmap
    k_fun.nMap = wires.n * pmap
    theta_fun.nMap = wires.n * pmap
    theta_fun.theta_rMap = wires.theta_r * pmap
    theta_fun.theta_sMap = wires.theta_s * pmap

    prob = Richards.RichardsProblem(
        M,
        hydraulic_conductivity=k_fun,
        water_retention=theta_fun,
        boundary_conditions=bc, initial_conditions=h,
        do_newton=False, method='mixed', debug=False
    )
    prob.timeSteps = [(5, 30, 1.2), (1200, 60)]

    # Create the survey
    locs = -np.arange(2, 38, 4.)/100
    locs = M.vectorCCx
    # times = np.arange(30, prob.timeMesh.vectorCCx[-1], 60)
    times = prob.timeMesh.vectorCCx[::4]
    if rx_type == 'saturation':
        rx = Richards.SaturationRx(locs, times)
    elif rx_type == 'pressure':
        rx = Richards.PressureRx(locs, times)
    survey = Richards.RichardsSurvey([rx])
    survey.pair(prob)

    Hs = prob.fields(model)
    data = survey.dpred(m=model, f=Hs)
    stdev = 0.01  # The standard deviation for the noise
    survey.makeSyntheticData(m=model, std=stdev, force=True, f=Hs)
    dmis = DataMisfit.l2_DataMisfit(survey)

    if plotIt:
        import matplotlib
        import matplotlib.pyplot as plt
        time_conversion = 60*60
        time_units = 'hours'

        plt.figure(figsize=(14, 9))
        plt.subplot(221)
        plt.plot(np.log10(prob.hydraulic_conductivity.Ks), M.gridCC)
        plt.title('True Model')
        plt.ylabel('Depth, m')
        plt.xlabel('Hydraulic conductivity, $log_{10}(K_s)$')
        plt.plot([-3.25]*len(locs), locs, 'ro')
        plt.legend(('True model', 'Data locations'))
        plt.subplot(222)
        plt.plot(times/time_conversion, data.reshape((-1, len(locs))))
        plt.title('True Data')
        plt.xlabel('Time, ' + time_units)
        plt.ylabel('Saturation')
        ax = plt.subplot(212)
        mesh2d = Mesh.TensorMesh(
            [prob.timeMesh.hx/time_conversion, prob.mesh.hx], '0N'
        )
        # sats = [theta_fun(_) for _ in Hs]
        clr = mesh2d.plotImage(np.c_[Hs][1:, :], ax=ax)
        cmap0 = matplotlib.cm.RdYlBu_r
        clr[0].set_cmap(cmap0)
        c = plt.colorbar(clr[0])
        c.set_label('Preasure Head $\psi$')
        plt.title('Preassure head, $\psi$, over time')
        plt.xlabel('Time, ' + time_units)
        plt.ylabel('Depth, m')
        plt.tight_layout()

    return prob, dmis, Hs


span_full = {
    'Ks': [1.5e-07, 6e-05],
    'alpha': [2.5, 13.5],
    'n': [1.1, 1.6],
    'theta_r': [0.01, 0.1],
    'theta_s': [0.3, 0.5]
}


ranges = [ksrange, trrange, tsrange, arange, nrange]
save_names = ['ks', 'tr', 'ts', 'a', 'n']
names = ['$K_s$', '$\\theta_r$', '$\\theta_s$', '$\\alpha$', '$n$']


def main(batch=None, rx_type=None):

    if batch is None:
        bcouples = couples
    else:
        bcouples = couples[(batch*2):(batch+1)*2]

    prob, dmis, Hs = get_problem(rx_type=rx_type)

    IMAGES = [np.zeros((n, m)) for _ in range(len(bcouples))]

    for ci, inds in enumerate(bcouples):
        irange = ranges[inds[0]]
        jrange = ranges[inds[1]]
        namei, namej = save_names[inds[0]], save_names[inds[1]]
        print('starting', namei, namej)
        for i in range(n):
            for j in range(m):
                mcopy = model.copy()
                mcopy[inds] = [irange[i], jrange[j]]
                IMAGES[ci][i, j] = dmis.eval(mcopy)
                print(ci, i, j)
        np.save(
            'image_{}_{}_{}_{}x{}'.format(rx_type, namei, namej, n, m),
            IMAGES[ci]
        )


def plot():
    import matplotlib
    import matplotlib.pyplot as plt
    fig, axes = plt.subplots(2, 5, figsize=(15, 8))
    for ci, inds in enumerate(couples):

        namei, namej = names[inds[0]], names[inds[1]]
        snamei, snamej = save_names[inds[0]], save_names[inds[1]]
        image = np.load(
            'image_{}_{}_{}_{}x{}.npy'.format(rx_type, snamei, snamej, n, m)
        )

        irange = ranges[inds[0]]
        jrange = ranges[inds[1]]
        X, Y = np.meshgrid(irange, jrange)

        ax = axes.flatten()[ci]
        clr = ax.contourf(X, Y, image.T, 40)
        cmap0 = matplotlib.cm.viridis_r
        clr.set_cmap(cmap0)
        if 'K_s' in namei:
            ax.set_xscale('log', nonposy='clip')

        ax.set_xlabel(namei)
        ax.set_ylabel(namej)
        ax.plot(model[inds[0]], model[inds[1]], 'r*')
    fig.tight_layout()
    plt.show()


if __name__ == '__main__':
    batch = None
    rx_types = ['pressure', 'saturation']
    rx_type = 'pressure'

    if len(sys.argv) > 2:
        rx_type = sys.argv[-2]
        assert rx_type in rx_types, 'choose pressure or saturation'
        batch = int(sys.argv[-1])
        assert 0 <= batch < 5, 'batch must be an int 0 to 4'
        print('batch', batch, 'rx_type', rx_type)

    main(batch=batch, rx_type=rx_type)
    if plotIt:
        plot()