def setUp_TDEM(prbtype="MagneticFluxDensity", rxcomp="bz", waveform="stepoff"):
    cs = 5.0
    ncx = 8
    ncy = 8
    ncz = 8
    npad = 4
    # hx = [(cs, ncx), (cs, npad, 1.3)]
    # hz = [(cs, npad, -1.3), (cs, ncy), (cs, npad, 1.3)]
    mesh = discretize.TensorMesh(
        [
            [(cs, npad, -1.3), (cs, ncx), (cs, npad, 1.3)],
            [(cs, npad, -1.3), (cs, ncy), (cs, npad, 1.3)],
            [(cs, npad, -1.3), (cs, ncz), (cs, npad, 1.3)],
        ],
        "CCC",
    )

    active = mesh.vectorCCz < 0.0
    activeMap = maps.InjectActiveCells(mesh, active, np.log(1e-8), nC=mesh.nCz)
    mapping = maps.ExpMap(mesh) * maps.SurjectVertical1D(mesh) * activeMap

    prb = getattr(tdem, "Simulation3D{}".format(prbtype))(mesh,
                                                          sigmaMap=mapping)

    rxtimes = np.logspace(-4, -3, 20)

    if waveform.upper() == "RAW":
        out = utils.VTEMFun(prb.times, 0.00595, 0.006, 100)
        wavefun = interp1d(prb.times, out)
        t0 = 0.006
        waveform = tdem.Src.RawWaveform(offTime=t0, waveFct=wavefun)
        prb.timeSteps = [(1e-3, 5), (1e-4, 5), (5e-5, 10), (5e-5, 10),
                         (1e-4, 10)]
        rxtimes = t0 + rxtimes

    else:
        waveform = tdem.Src.StepOffWaveform()
        prb.timeSteps = [(1e-05, 10), (5e-05, 10), (2.5e-4, 10)]

    rxOffset = 10.0
    rx = getattr(tdem.Rx, "Point{}".format(rxcomp[:-1]))(np.r_[rxOffset, 0.0,
                                                               -1e-2], rxtimes,
                                                         rxcomp[-1])
    src = tdem.Src.MagDipole([rx],
                             loc=np.array([0.0, 0.0, 0.0]),
                             waveform=waveform)

    survey = tdem.Survey([src])

    prb.Solver = Solver

    m = np.log(1e-1) * np.ones(prb.sigmaMap.nP) + 1e-2 * np.random.rand(
        prb.sigmaMap.nP)

    prb.pair(survey)
    mesh = mesh

    return prb, m, mesh
def get_survey(prob, t0):

    out = utils.VTEMFun(prob.times, 0.00595, 0.006, 100)
    wavefun = interp1d(prob.times, out)

    waveform = tdem.Src.RawWaveform(offTime=t0, waveFct=wavefun)
    src = tdem.Src.MagDipole([], waveform=waveform, loc=np.array([0.0, 0.0, 0.0]))

    return tdem.Survey([src])
def run(plotIt=True):

    cs, ncx, ncz, npad = 5.0, 25, 24, 15
    hx = [(cs, ncx), (cs, npad, 1.3)]
    hz = [(cs, npad, -1.3), (cs, ncz), (cs, npad, 1.3)]
    mesh = discretize.CylMesh([hx, 1, hz], "00C")

    active = mesh.vectorCCz < 0.0
    layer = (mesh.vectorCCz < -50.0) & (mesh.vectorCCz >= -150.0)
    actMap = maps.InjectActiveCells(mesh, active, np.log(1e-8), nC=mesh.nCz)
    mapping = maps.ExpMap(mesh) * maps.SurjectVertical1D(mesh) * actMap
    sig_half = 1e-3
    sig_air = 1e-8
    sig_layer = 1e-2
    sigma = np.ones(mesh.nCz) * sig_air
    sigma[active] = sig_half
    sigma[layer] = sig_layer
    mtrue = np.log(sigma[active])

    x = np.r_[30, 50, 70, 90]
    rxloc = np.c_[x, x * 0.0, np.zeros_like(x)]

    prb = TDEM.Simulation3DMagneticFluxDensity(mesh,
                                               sigmaMap=mapping,
                                               solver=Solver)
    prb.time_steps = [
        (1e-3, 5),
        (1e-4, 5),
        (5e-5, 10),
        (5e-5, 5),
        (1e-4, 10),
        (5e-4, 10),
    ]
    # Use VTEM waveform
    out = EMutils.VTEMFun(prb.times, 0.00595, 0.006, 100)

    # Forming function handle for waveform using 1D linear interpolation
    wavefun = interp1d(prb.times, out)
    t0 = 0.006
    waveform = TDEM.Src.RawWaveform(offTime=t0, waveFct=wavefun)

    rx = TDEM.Rx.PointMagneticFluxTimeDerivative(
        rxloc,
        np.logspace(-4, -2.5, 11) + t0, "z")
    src = TDEM.Src.CircularLoop([rx],
                                waveform=waveform,
                                loc=np.array([0.0, 0.0, 0.0]),
                                radius=10.0)
    survey = TDEM.Survey([src])
    prb.survey = survey

    # create observed data
    data = prb.make_synthetic_data(mtrue,
                                   relative_error=0.02,
                                   noise_floor=1e-11)

    dmisfit = data_misfit.L2DataMisfit(simulation=prb, data=data)
    regMesh = discretize.TensorMesh([mesh.hz[mapping.maps[-1].indActive]])
    reg = regularization.Simple(regMesh)
    opt = optimization.InexactGaussNewton(maxIter=5, LSshorten=0.5)
    invProb = inverse_problem.BaseInvProblem(dmisfit, reg, opt)
    target = directives.TargetMisfit()
    # Create an inversion object
    beta = directives.BetaSchedule(coolingFactor=1.0, coolingRate=2.0)
    betaest = directives.BetaEstimate_ByEig(beta0_ratio=1e0)
    invProb.beta = 1e2
    inv = inversion.BaseInversion(invProb, directiveList=[beta, target])
    m0 = np.log(np.ones(mtrue.size) * sig_half)
    prb.counter = opt.counter = utils.Counter()
    opt.remember("xc")
    mopt = inv.run(m0)

    if plotIt:
        fig, ax = plt.subplots(1, 2, figsize=(10, 6))
        Dobs = data.dobs.reshape((len(rx.times), len(x)))
        Dpred = invProb.dpred.reshape((len(rx.times), len(x)))
        for i in range(len(x)):
            ax[0].loglog(rx.times - t0, -Dobs[:, i].flatten(), "k")
            ax[0].loglog(rx.times - t0, -Dpred[:, i].flatten(), "k.")
            if i == 0:
                ax[0].legend(("$d^{obs}$", "$d^{pred}$"), fontsize=16)
        ax[0].set_xlabel("Time (s)", fontsize=14)
        ax[0].set_ylabel("$db_z / dt$ (nT/s)", fontsize=16)
        ax[0].set_xlabel("Time (s)", fontsize=14)
        ax[0].grid(color="k", alpha=0.5, linestyle="dashed", linewidth=0.5)

        plt.semilogx(sigma[active], mesh.vectorCCz[active])
        plt.semilogx(np.exp(mopt), mesh.vectorCCz[active])
        ax[1].set_ylim(-600, 0)
        ax[1].set_xlim(1e-4, 1e-1)
        ax[1].set_xlabel("Conductivity (S/m)", fontsize=14)
        ax[1].set_ylabel("Depth (m)", fontsize=14)
        ax[1].grid(color="k", alpha=0.5, linestyle="dashed", linewidth=0.5)
        plt.legend(["$\sigma_{true}$", "$\sigma_{pred}$"])
def halfSpaceProblemAnaDiff(
    meshType,
    srctype="MagDipole",
    sig_half=1e-2,
    rxOffset=50.0,
    bounds=None,
    plotIt=False,
    rxType="MagneticFluxDensityz",
):

    if bounds is None:
        bounds = [1e-5, 1e-3]
    if meshType == "CYL":
        cs, ncx, ncz, npad = 15.0, 30, 10, 15
        hx = [(cs, ncx), (cs, npad, 1.3)]
        hz = [(cs, npad, -1.3), (cs, ncz), (cs, npad, 1.3)]
        mesh = discretize.CylMesh([hx, 1, hz], "00C")

    elif meshType == "TENSOR":
        cs, nc, npad = 20.0, 13, 5
        hx = [(cs, npad, -1.3), (cs, nc), (cs, npad, 1.3)]
        hy = [(cs, npad, -1.3), (cs, nc), (cs, npad, 1.3)]
        hz = [(cs, npad, -1.3), (cs, nc), (cs, npad, 1.3)]
        mesh = discretize.TensorMesh([hx, hy, hz], "CCC")

    active = mesh.vectorCCz < 0.0
    actMap = maps.InjectActiveCells(mesh, active, np.log(1e-8), nC=mesh.nCz)
    mapping = maps.ExpMap(mesh) * maps.SurjectVertical1D(mesh) * actMap

    prb = tdem.Simulation3DMagneticFluxDensity(mesh, sigmaMap=mapping)
    prb.Solver = Solver
    prb.timeSteps = [(1e-3, 5), (1e-4, 5), (5e-5, 10), (5e-5, 10), (1e-4, 10)]
    out = utils.VTEMFun(prb.times, 0.00595, 0.006, 100)
    wavefun = interp1d(prb.times, out)
    t0 = 0.006
    waveform = tdem.Src.RawWaveform(offTime=t0, waveFct=wavefun)

    rx = getattr(tdem.Rx,
                 "Point{}".format(rxType[:-1]))(np.array([[rxOffset, 0.0,
                                                           0.0]]),
                                                np.logspace(-4, -3, 31) + t0,
                                                rxType[-1])

    if srctype == "MagDipole":
        src = tdem.Src.MagDipole([rx],
                                 waveform=waveform,
                                 loc=np.array([0, 0.0, 0.0]))
    elif srctype == "CircularLoop":
        src = tdem.Src.CircularLoop([rx],
                                    waveform=waveform,
                                    loc=np.array([0.0, 0.0, 0.0]),
                                    radius=13.0)

    survey = tdem.Survey([src])
    prb.pair(survey)

    sigma = np.ones(mesh.nCz) * 1e-8
    sigma[active] = sig_half
    sigma = np.log(sigma[active])

    if srctype == "MagDipole":
        bz_ana = mu_0 * analytics.hzAnalyticDipoleT(rx.locations[0][0] + 1e-3,
                                                    rx.times - t0, sig_half)
    elif srctype == "CircularLoop":
        bz_ana = mu_0 * analytics.hzAnalyticCentLoopT(13, rx.times - t0,
                                                      sig_half)

    bz_calc = prb.dpred(sigma)
    ind = np.logical_and(rx.times - t0 > bounds[0], rx.times - t0 < bounds[1])
    log10diff = np.linalg.norm(
        np.log10(np.abs(bz_calc[ind])) -
        np.log10(np.abs(bz_ana[ind]))) / np.linalg.norm(
            np.log10(np.abs(bz_ana[ind])))

    print(" |bz_ana| = {ana} |bz_num| = {num} |bz_ana-bz_num| = {diff}".format(
        ana=np.linalg.norm(bz_ana),
        num=np.linalg.norm(bz_calc),
        diff=np.linalg.norm(bz_ana - bz_calc),
    ))
    print("Difference: {}".format(log10diff))

    if plotIt is True:
        plt.loglog(
            rx.times[bz_calc > 0] - t0,
            bz_calc[bz_calc > 0],
            "r",
            rx.times[bz_calc < 0] - t0,
            -bz_calc[bz_calc < 0],
            "r--",
        )
        plt.loglog(rx.times - t0, abs(bz_ana), "b*")
        plt.title("sig_half = {:e}".format(sig_half))
        plt.show()

    return log10diff