<jupyter_start><jupyter_text>set up a simple test example<jupyter_code>def waveform(t, t_peak=2e-3, width=20, amplitude=1):
    t = np.log10(t)
    t_peak = np.log10(t_peak)
    width = np.log10(width)
    return amplitude * np.exp(-(t - t_peak)**2/(2*width**2))

def sigmoid(x, x0=0, slope=1):
    return np.arctan(slope * (x-x0))/np.pi + 0.5

def depth_distribution(z, t, dz=200, slope=1e-1):
    Z = np.outer(z, np.ones(len(t)))
    top = model.casing_z.max()
    time_fact = -np.log(t + t[1])
    time_fact = -10*(t)/np.max(t)
    T = np.outer(np.ones(len(z)), time_fact)
    bottom = -(model.casing_z.min()*T)-dz 
    return sigmoid(Z, bottom, slope) * sigmoid(-Z, -top, slope)

def source_se(mesh, t, z, dz=200, slope=None):
    sm = np.zeros(mesh.nE)
    wave = np.outer(np.ones(len(z)), waveform(t))
    sm = depth_distribution(z, t, dz=dz, slope=slope) * waveform(t)
    
    return sm
# z = np.load(f"{trial_directory}/z_currents.npy")
csx = mesh.hx.min()

# xinds = (mesh.gridFz[:, 0] < model.casing_b + csx/2) & (mesh.gridFz[:, 0] > model.casing_b - csx/2)

# # xinds = (mesh.gridFz[:, 0] < model.casing_b/2 + csx/2) & (mesh.gridFz[:, 0] > model.casing_b/2 - csx/2)
# zinds = (mesh.gridFz[:, 2] >= model.casing_z.min()) & (mesh.gridFz[:, 2] <= model.casing_z.max())
# src_inds_bool = xinds & zinds
# src_inds = np.where(src_inds_bool)[0]
# z = mesh.gridFz[src_inds, 2]

# print(np.unique(mesh.gridFz[xinds, 0]))
z = mesh.vectorNz[(mesh.vectorNz >= model.casing_z.min()) & (mesh.vectorNz <= model.casing_z.max())]
src_vec = -1*source_se(mesh, sim.times, z, dz=300, slope=1e-2)
fig, ax = plt.subplots(1, 1)
plt.colorbar(ax.pcolormesh(sim.times, z, src_vec), ax=ax)
ax.set_xscale("log")
ax.set_xlim(1e-6, sim.times.max())
ax.set_xlabel("time (s)")
ax.set_ylabel("z")
mesh = discretize.CylMesh([mesh3D.hx, 1, mesh3D.hz], x0=mesh3D.x0)
sim = tdem.Problem3D_j(mesh=mesh, time_steps=sim3D.time_steps, solver=Solver, solver_opts=solver_opts, sigma=model.sigma(mesh))
class SourceMap(maps.IdentityMap):
    def __init__(self, mesh, spatial_index, nT):
        spatial_inds = np.zeros(mesh.nF, dtype=bool)
        spatial_inds[spatial_index] = True
        time_index = np.ones(nT, dtype=bool)
        index = np.kron(time_index, spatial_inds)
        
        super(SourceMap, self).__init__(mesh=mesh, nP=index.sum())
        self.index = index
        self.nF = self.mesh.nF
        self.nT = nT
    
    def _transform(self, m):
        src = np.zeros(self.nF*self.nT)
        src[self.index] = m
        return src
    
    def deriv(self, m, v=None):
        if v is None:
            P = sp.csr_matrix(
                (np.ones(self.nP), (self.index.nonzero()[0], range(self.nP))), 
                shape=(self.nF*self.nT, self.nP)
            )
            return P
        return self._transform(v)
                  
# xinds = (mesh.gridFz[:, 0] < model.casing_b + csx/2) & (mesh.gridFz[:, 0] > model.casing_b - csx/2)

xinds = (mesh.gridFz[:, 0] < csx) #& (mesh.gridFz[:, 0] > model.casing_b/2 - csx/2)
zinds = (mesh.gridFz[:, 2] >= model.casing_z.min()) & (mesh.gridFz[:, 2] <= model.casing_z.max())
src_inds_bool = xinds & zinds
src_inds = np.where(src_inds_bool)[0]
src_inds = src_inds + mesh.nFx

# z = mesh.gridFz[src_inds, 2]
src_map = SourceMap(mesh, src_inds, sim.nT+1)
# src_map * discretize.utils.mkvc(src_vec)
src_map.test(discretize.utils.mkvc(src_vec))
# fig, ax = plt.subplots(1, 1)
# tind = 40
# plotme = (src_map * discretize.utils.mkvc(src_vec)).reshape((mesh.nEy, sim.nT+1), order="F")[:, tind]

# plt.colorbar(
#     mesh.plotImage(
#         mesh.aveEy2CC * plotme, ax=ax, mirror=True, 
#         pcolorOpts={"cmap":"BrBG_r"}, clim=plotme.max()*np.r_[-1, 1]
#     )[0], ax=ax
# )
# ax.set_xlim(0.2*np.r_[-1, 1])
# ax.set_ylim([-2000, 10])
class MySimulation(tdem.Problem3D_j):
    
    source, source_map, source_deriv = props.Invertible("source term")
    
    j_inds = None
    
    survey = None
    
    def fields(self, m):
        self.model = m
        
#         source = self.source.reshape(sim.mesh.nEy, sim.nT+1, order="F")
        f = np.zeros((self.mesh.nF, self.nT+1))

        # this assumes the initial condition is zero. 

        # timestep to solve forward
        Ainv = None
        
        for tInd, dt in enumerate(self.timeSteps):
            # keep factors if dt is the same as previous step b/c A will be the
            # same
            if Ainv is not None and (
                tInd > 0 and abs(dt-self.timeSteps[tInd - 1]) >
                self.dt_threshold
            ):
                Ainv.clean()
                Ainv = None

            if Ainv is None:
                A = self.getAdiag(tInd)
                Ainv = self.solver(A)

            rhs = self.getRHS(tInd+1)  # this is on the nodes of the time mesh
            Asubdiag = self.getAsubdiag(tInd)

            # taking a step
            sol = Ainv * (rhs - Asubdiag * f[:, tInd])
            f[:, tInd+1] = sol

        # clean factors and return
        Ainv.clean()
        return f

    def dpred(self, m, f=None):
        if f is None: 
            f = self.fields(m)
        return mkvc(f)[self.j_inds]

    def getRHS(self, tInd):
        src = self.source.reshape((self.mesh.nF, self.nT+1), order="F")
        rhs = -1/self.time_steps[tInd-1] * (src[:, tInd] - src[:, tInd-1])
        if self._makeASymmetric:
            return self.MfRho.T * rhs
        return rhs

    def getRHSDeriv(self, tInd, v, adjoint=False):

        if adjoint:
            if self._makeASymmetric:
                v = self.MfRho * v
            else:
                v = v 
            
            rhs = np.zeros((self.mesh.nF, self.nT+1))
            rhs[:, tInd] = v
            rhs[:, tInd-1] = -v
            return -1/self.time_steps[tInd-1] * rhs

        v = v.reshape((self.mesh.nF, self.nT+1), order="F")
        rhs = -1/self.time_steps[tInd-1] * (v[:, tInd] - v[:, tInd-1])

        if self._makeASymmetric:
            return self.MfRho.T * rhs
        return rhs
        
    def Jvec(self, m, v, f=None):
        self.model = m
        v = self.source_deriv * v
#         v = v.reshape((self.mesh.nF, self.nT+1), order="F")
        df_dm_v = np.zeros((self.mesh.nF, self.nT+1))

        # timestep to solve forward
        Ainv = None
        for tInd, dt in enumerate(self.timeSteps):
            # keep factors if dt is the same as previous step b/c A will be the
            # same
            if Ainv is not None and (
                tInd > 0 and abs(dt-self.timeSteps[tInd - 1]) > self.dt_threshold
            ):
                Ainv.clean()
                Ainv = None

            if Ainv is None:
                A = self.getAdiag(tInd)
                Ainv = self.solver(A)

            rhs_deriv = self.getRHSDeriv(tInd+1, v)  # this is on the nodes of the time mesh
            Asubdiag = self.getAsubdiag(tInd)

            # taking a step
            sol = Ainv * (rhs_deriv - Asubdiag * df_dm_v[:, tInd])
            df_dm_v[:, tInd+1] = sol

        # clean factors and return
        Ainv.clean()
        
        return mkvc(df_dm_v)[self.j_inds]
        
    def Jtvec(self, m, v, f=None):
        self.model = m
        if self.j_inds is not None:
            v_tmp = np.zeros(self.mesh.nF*(self.nT+1))
            v_tmp[self.j_inds] = v
            v = v_tmp.reshape(self.mesh.nF, self.nT+1, order="F")
        
        df_dmT_v = np.zeros((self.mesh.nF, self.nT+1))  # the source is defined on faces
        
        # timestep to solve forward
        ATinv = None
        for tInd in reversed(range(self.nT)):
            dt = self.time_steps[tInd]
            # keep factors if dt is the same as previous step b/c A will be the
            # same
            if ATinv is not None and (
                tInd <= self.nT and abs(dt-self.timeSteps[tInd + 1]) > self.dt_threshold
            ):
                ATinv.clean()
                ATinv = None

            if ATinv is None:
                AT = self.getAdiag(tInd).T
                ATinv = self.solver(AT)    

    #         ATinv_v = ATinv * v[:, tInd+1]
            if tInd < self.nT - 1:
                AsubdiagT = self.getAsubdiag(tInd+1).T
                sol = ATinv * (v[:, tInd+1] - AsubdiagT * sol)
            else: 
                sol = ATinv * v[:, tInd+1]
            rhs_deriv = self.getRHSDeriv(tInd+1, sol, adjoint=True)  # this is on the nodes of the time mesh

            df_dmT_v += rhs_deriv

        # clean factors and return
        ATinv.clean()
        return self.source_deriv.T * mkvc(df_dmT_v)
src_map = SourceMap(mesh, src_inds, sim.nT+1)
x_bounds = np.r_[10, 1000]
z_bounds = np.r_[-1200, 0.]
j_inds = get_j_inds(
    mesh, nT=sim.nT, x_bounds=x_bounds, z_bounds=z_bounds
)

mysim = MySimulation(
    mesh=mesh, time_steps=sim3D.time_steps, 
    solver=Solver, solver_opts=solver_opts, sigma=model.sigma(mesh),
    source_map=src_map, j_inds=j_inds
)
s0 = mkvc(src_vec)

def derChk(m):
    return [mysim.dpred(m), lambda mx: mysim.Jvec(m, mx)]

discretize.Tests.checkDerivative(derChk, s0, plotIt=False, num=3, eps=1e-9)
v = np.random.rand(len(s0))
w = np.random.rand(j_inds.sum())

wTJv = w.dot(mysim.Jvec(s0, v))
vTJTw = v.dot(mysim.Jtvec(s0, w))

print(f"{wTJv:1.2e}, {vTJTw:1.2e}, {np.abs(wTJv - vTJTw):1.2e}")<jupyter_output>-2.31e-05, -2.31e-05, 7.59e-13<jupyter_text>set up an inversion<jupyter_code>j_obs = mysim.dpred(s0)
plt.hist(np.log10(np.abs(j_obs+1e-20)), 20);
std = 0.05
floor = 1e-12
mysurvey = survey.BaseSurvey()
mysurvey._vnD = np.r_[len(j_obs)]
mysim.survey = mysurvey
mydata = data.Data(mysurvey, dobs=j_obs, standard_deviation=std, noise_floor=floor)
dmis = data_misfit.L2DataMisfit(data=mydata, simulation=mysim)
dz = np.diff(z)
dz = np.hstack([dz, np.r_[dz[-1]]])

dt = mysim.time_mesh.hx
dt = np.hstack([dt, np.r_[dt[-1]]])

reg_mesh = discretize.TensorMesh([dz, dt], x0=np.r_[z.min(), mysim.times[0]])
reg = regularization.Tikhonov(
    mesh=reg_mesh, 
    alpha_x=dz.mean()**2, # space
    alpha_y=dt.mean()**2, # time
)
opt = optimization.InexactGaussNewton(maxIter=20)
invprob = inverse_problem.BaseInvProblem(dmis, reg, opt)

directiveslist = [
    directives.BetaEstimate_ByEig(beta0_ratio=1e2),
    directives.BetaSchedule(),
    directives.TargetMisfit(),
]

inv = inversion.BaseInversion(invprob, directiveList=directiveslist)
print(f"{mysurvey.nD:1.2e}, {0.5*mysurvey.nD:1.2e}")
m0 = 0 * s0 
mrec = inv.run(m0)
fig, ax = plt.subplots(1, 2, figsize=(12, 4))
plt.colorbar(ax[0].pcolormesh(sim.times, z, mrec.reshape((len(z), sim.nT+1), order="F")), ax=ax[0])
ax[0].set_title("recovered")

plt.colorbar(ax[1].pcolormesh(sim.times, z, s0.reshape((len(z), sim.nT+1), order="F")), ax=ax[1])
ax[1].set_title("true")

for a in ax: 
    a.set_xscale("log")
    a.set_xlim(1e-6, sim.times.max())
    a.set_xlabel("time (s)")
    a.set_ylabel("z")
j_pred = mysim.fields(mrec)
j_true = mysim.fields(s0)
tind = 40
fig, ax = plt.subplots(1, 2, figsize=(12, 4))

xlim = np.r_[-1000, 1000]
zlim = np.r_[-1500, 100]
sample_grid = np.r_[5., 5.]
clim = np.r_[1e-11, 1e-4]

out = mesh.plotImage(
    mesh.aveF2CCV * j_pred[:, tind], 
    view="vec",
    vType="CCv",
    ax=ax[0], mirror=True,
    range_x=xlim, 
    range_y=zlim,
    sample_grid = sample_grid,
    pcolorOpts={"norm":LogNorm()},
    clim=clim, 
    stream_threshold=clim.min()
)

plt.colorbar(out[0], ax=ax[0])

ax[0].set_title(f"recovered, t={sim.times[tind]*1e3:1.1e}ms")


out2 = mesh.plotImage(
    mesh.aveF2CCV * j_true[:, tind], 
    view="vec",
    vType="CCv",
    ax=ax[1], mirror=True,
    range_x=xlim, 
    range_y=zlim,
    sample_grid = sample_grid,
    pcolorOpts={"norm":LogNorm()},
    clim=clim, 
    stream_threshold=clim.min()
)
plt.colorbar(out2[0], ax=ax[1])
# ax[1].set_ylim(np.r_[-max_depth, top])
ax[1].set_ylabel('z (m)')
ax[1].set_title(f"true , t={sim.times[tind]*1e3:1.1e}ms")

for a in ax:
    a.set_aspect(1)<jupyter_output><empty_output><jupyter_text>invert the simulated data<jupyter_code>model_dict, fields_dict, sim3D, mesh3D = load_trial(trial_directory)
jd_3D = fields_dict["casing"][:, "j", :] - fields_dict["approx_casing"][:, "j", :]
jd_x = (jd_3D[:mesh3D.nFx, :]).reshape(np.hstack([mesh3D.vnFx, np.r_[sim.nT+1]]), order="F")
jd_z = (jd_3D[mesh3D.vnF[:2].sum():, :]).reshape(np.hstack([mesh3D.vnFz, np.r_[sim.nT+1]]), order="F")

# grab a slice through theta
theta_ind = 3
jd = np.hstack([
    discretize.utils.mkvc(jd_x[:, theta_ind, :, :]),
    discretize.utils.mkvc(jd_z[:, theta_ind, :, :]),
])

jd_obs = jd[j_inds]
plt.hist(np.log10(np.abs(jd+1e-30)), 20);
model_halfspace = model.copy()
model_halfspace.sigma_casing = model.sigma_back
model_halfspace.sigma_inside = model.sigma_back
std = 0.05
floor = 1e-9
mysurvey = survey.BaseSurvey()
mysurvey._vnD = np.r_[len(jd_obs)]

# sigma = model.sigma_air * np.ones(mesh.nC)
# sigma[mesh.gridCC[:, 2] <0] = model.sigma_back
mysim.sigma = model.sigma(mesh)
mysim.survey = mysurvey
mydata2 = data.Data(mysurvey, dobs=jd_obs, standard_deviation=std, noise_floor=floor)

np.unique(mysim.sigma)
dmis = data_misfit.L2DataMisfit(data=mydata2, simulation=mysim)

reg = regularization.Tikhonov(
    mesh=reg_mesh, 
    alpha_x=dz.mean()**2, # space
    alpha_y=dt.mean()**2/100, # time
)
opt = optimization.InexactGaussNewton(maxIter=50)
invprob = inverse_problem.BaseInvProblem(dmis, reg, opt)

directiveslist = [
    directives.BetaEstimate_ByEig(beta0_ratio=1e2),
    directives.BetaSchedule(coolingFactor=4.),
    directives.TargetMisfit(),
    directives.SaveOutputDictEveryIteration(),
]

inv = inversion.BaseInversion(invprob, directiveList=directiveslist)
print(f"{mysurvey.nD:1.2e}, {0.5*mysurvey.nD:1.2e}")
# m0 = np.zeros_like(s0)
srec = inv.run(srec)<jupyter_output>SimPEG.InvProblem will set Regularization.mref to m0.

    SimPEG.InvProblem is setting bfgsH0 to the inverse of the eval2Deriv.
    ***Done using same Solver and solverOpts as the problem***<jupyter_text>todo: look at the regularization<jupyter_code>reg.objfcts
for multiplier, fct in zip(reg.multipliers, reg.objfcts):
    print(f"{fct.__class__.__name__}: {multiplier * fct(srec)}")
fig, ax = plt.subplots(1, 1)
plt.colorbar(ax.pcolormesh(sim.times, z, srec.reshape((len(z), sim.nT+1), order="F")), ax=ax)
ax.set_title("recovered")

ax.set_xscale("log")
ax.set_xlim(1e-6, sim.times.max())
ax.set_xlabel("time (s)")
ax.set_ylabel("z")
fig, ax = plt.subplots(1, 1)

cm = plt.get_cmap('viridis')
c_norm = LogNorm(vmin=sim.time_mesh.vectorCCx[0], vmax=sim.time_mesh.vectorCCx[-1])
scalar_map = cmap.ScalarMappable(norm=c_norm, cmap=cm)
scalar_map.set_array([])

plotme = srec.reshape((len(z), sim.nT+1), order="F")

for time_ind in range(sim.nT)[::int(sim.nT/20)]:
    color = scalar_map.to_rgba(sim.time_mesh.vectorCCx[time_ind])
    ax.plot(z, plotme[:, time_ind], color=color)

ax.set_xlim(z.max(), z.min())

cbar_ax = fig.add_axes([1, 0.1, 0.02, 0.8])
cb = plt.colorbar(scalar_map, cbar_ax)
cb.set_label('time (s)')
cb.ax.invert_yaxis()
fig, ax = plt.subplots(1, 1)

cm = plt.get_cmap('viridis')
c_norm = Normalize(vmin=z.min(), vmax=z.max())
scalar_map = cmap.ScalarMappable(norm=c_norm, cmap=cm)
scalar_map.set_array([])

plotme = srec.reshape((len(z), sim.nT+1), order="F")

for z_ind in range(len(z))[::int(len(z)/20)]:
    color = scalar_map.to_rgba(z[z_ind])
    ax.semilogx(sim.time_mesh.vectorNx, plotme[z_ind, :], color=color)

# ax.set_xlim(z.max(), z.min())

cbar_ax = fig.add_axes([1, 0.1, 0.02, 0.8])
cb = plt.colorbar(scalar_map, cbar_ax)
cb.set_label('z (m)')
cb.ax.invert_yaxis()
j_rec = mysim.fields(srec)
model_names = list(model_dict.keys())
# sim.survey.source_list = sim.survey.source_list
viewer = casing_sim.FieldsViewer(
    mesh=mesh3D, model_parameters_dict=model_dict, survey_dict={key: sim3D.survey for key in model_names}, 
    fields_dict=fields_dict, model_keys=model_names, primary_key="casing"
)
tind = 90
fig, ax = plt.subplots(1, 2, figsize=(12, 4))

xlim = 1000*np.r_[-1, 1]
zlim = np.r_[-1500, 100]
sample_grid = np.r_[xlim.max()/50, 5.]
clim = np.r_[1e-11, 1e-6]

out = mesh.plotImage(
    mesh.aveF2CCV * j_rec[:, tind], 
    view="vec",
    vType="CCv",
    ax=ax[0], mirror=True,
    range_x=xlim, 
    range_y=zlim,
    sample_grid = sample_grid,
    pcolorOpts={"norm":LogNorm()},
    clim=clim, 
    stream_threshold=clim.min()
)

plt.colorbar(out[0], ax=ax[0])

ax[0].set_title(f"recovered, t={sim.times[tind]*1e3:1.1e}ms")


out2 = viewer.plot_cross_section(
    ax=ax[1], clim=clim, zlim=zlim,
    xlim=xlim,
    view='j', theta_ind=3, time_ind=tind, 
    model_key='approx_casing', show_cb=True, casing_outline=False,
    prim_sec="secondary"
#     stream_opts={"density":0.75, "color": "k", "arrowsize": 2}
)
# ax[1].set_ylim(np.r_[-max_depth, top])
ax[1].set_ylabel('z (m)')
ax[1].set_title(f"true , t={sim.times[tind]*1e3:1.1e}ms")

for a in ax:
    a.set_aspect(1)
viewer.widget_cross_section()
mesh3D
def compute_sm(tind, key="approx_casing", primary_key="casing"):
    j_diff = fields_dict[primary_key][:, 'j', tind] - fields_dict[key][:, 'j', tind]
    ja = fields_dict[key][:, 'j', tind]
    
    C = mesh3D.edgeCurl
    MfRhod = (
        mesh3D.getFaceInnerProduct(1./model_dict[primary_key].sigma(mesh3D) - 1./model_dict[key].sigma(mesh3D)) 
    )
    
    return - C.T * MfRhod * (ja + jd)
# def compute_sm_error(tind, key="approx_casing", primary_key="casing", separate=False):
#     j_diff =  fields_dict[key][:, 'j', tind] - fields_dict[primary_key][:, 'j', tind]
#     dhdt_diff = fields_dict[key][:, 'dhdt', tind] - fields_dict[primary_key][:, 'dhdt', tind]
    
#     C = mesh3D.edgeCurl
#     MfRho = mesh3D.getFaceInnerProduct(1./model_dict[key].sigma(mesh3D))
#     MeMu = mesh3D.getEdgeInnerProduct(model_dict[key].mur(mesh3D)*mu_0)
    
#     if separate is True:
#         return (C.T * MfRho * j_diff, MeMu * dhdt_diff)
#     return C.T * MfRho * j_diff + MeMu * dhdt_diff
def theta_index_mirror(theta_ind):
    return (
        theta_ind+int(mesh3D.vnC[1]/2)
        if theta_ind < int(mesh3D.vnC[1]/2)
        else theta_ind-int(mesh3D.vnC[1]/2)
    )
def plot_sm_error(tind, theta_ind=0, key="approx_casing", primary_key="casing", ax=None, clim=None, xlim=None, zlim=None):
    if ax is None:
        fig, ax = plt.subplots(1, 1)
        
    plotme = compute_sm_error(tind, key, primary_key)

    if len(mesh.hy) == 1:
        plotme = mesh3D.aveE2CC * plotme
    else:
        plotme = (mesh3D.aveE2CCV * plotme)[mesh3D.nC:2*mesh3D.nC]
    
    theta_ind_mirror = theta_index_mirror(theta_ind)
    plotme = plotme.reshape(mesh3D.vnC, order="F")
    mirror_data = discretize.utils.mkvc(-plotme[:, theta_ind_mirror, :])
    plotme = discretize.utils.mkvc(plotme[:, theta_ind, :])

    norm = SymLogNorm(
        clim[0] if clim is not None else
        np.max([1e-20, np.min(np.absolute(plotme))])
    )
    clim = clim[1]*np.r_[-1., 1.] if clim is not None else None
    
    mesh2D = discretize.CylMesh([mesh.hx, 1, mesh.hz], x0=mesh.x0)
    
    out = mesh2D.plotImage(
        plotme, ax=ax, pcolorOpts = {'norm': norm, "cmap": "BrBG_r"}, clim=clim,
        mirror_data=mirror_data, mirror=True
    )

    cb = plt.colorbar(out[0], ax=ax)
    
    ax.set_xlim(xlim)
    ax.set_ylim(zlim)
    

def widget_sm(tind, theta_ind, clim_min, clim_max):
    fig, ax = plt.subplots(1, 1, figsize=(4, 6))
    plot_sm_error(
        tind=tind, theta_ind=theta_ind, ax=ax, xlim=0.15*np.r_[-1, 1], zlim=[-1.1*model_dict["casing"].casing_l, 10],
        clim=np.r_[clim_min, clim_max],
    )

ipywidgets.interact(
    widget_sm, 
    tind = ipywidgets.IntSlider(min=0, max=sim.nT+1, value=0),
    theta_ind = ipywidgets.IntSlider(min=0, max=len(mesh3D.hy), value=3),
    clim_min = ipywidgets.FloatText(value=1e-6),
    clim_max = ipywidgets.FloatText(value=1e-4),
#     term = ipywidgets.ToggleButtons(options=[0, 1, 2], value=2)
)
sm = discretize.utils.mkvc(compute_sm_error(tind=40, key="approx_casing", primary_key="casing"))
smx = sm[:mesh3D.vnE[0]]
smy = sm[mesh3D.vnE[0]:mesh3D.vnE[:2].sum()]
smz = sm[mesh3D.vnE[:2].sum():]
fig, ax = plt.subplots(1, 3, figsize=(12,4))
ylim = np.r_[1e-7, 1e-3]

for a, s, title in zip(ax, [smx, smy, smz], ["x", "y", "z"]):
    a.semilogy(np.abs(s))
    a.set_title(title)
    a.set_ylim(ylim)
inds = np.abs(smy) > 1e-6
source_grid = mesh3D.gridEy[inds, :]
np.unique(source_grid[:, 0])
SMY = []

for t in range(sim.nT+1):
    smi = discretize.utils.mkvc(compute_sm_error(tind=t, key="approx_casing", primary_key="casing"))
    smyi = discretize.utils.mkvc(smi[mesh3D.vnE[0]:mesh3D.vnE[:2].sum()])
    SMY.append(smyi)
SMY = np.vstack(SMY).T
radii = np.r_[0.02, 0.025, 0.03, 0.035, 0.04]
eps = 1e-3
smy_casing = []

for r in radii: 
    inds = (
        (mesh3D.gridEy[:, 0] >= r-eps) & (mesh3D.gridEy[:, 0] <= r+eps) &
        (mesh3D.gridEy[:, 1] == mesh3D.vectorCCy[3]) & 
        (mesh3D.gridEy[:, 2] >= -model.casing_l) & (mesh3D.gridEy[:, 2] <= 0)
    )
    smy_casing.append(SMY[inds, :])

smy_casing[0].shape
fig, ax = plt.subplots(1, len(radii), figsize=(5*len(radii), 4))

for i in range(len(smy_casing)):
    plt.colorbar(
        ax[i].pcolormesh(sim.times, z, (smy_casing[i]).reshape((len(z), sim.nT+1), order="F")), 
        ax=ax[i]
    )
    ax[i].set_title(f"r = {radii[i]}")

for a in ax: 
    a.set_xscale("log")
    a.set_xlim(1e-6, sim.times.max())
    a.set_xlabel("time (s)")
    a.set_ylabel("z")
plt.tight_layout()
tind = 50
fig, ax = plt.subplots(1, 1) 

for i in range(len(smy_casing)):
    ax.plot(z, smy_casing[i][:, tind], label=f"r = {radii[i]}")
    
ax.legend()
ax.ticklabel_format(axis="y", scilimits=(0, 0))
ax.set_title(f"time {1e3*sim.times[tind]:1.1e}ms")
ax.set_xlim(ax.get_xlim()[::-1])
ax.set_xlabel("z (m)")<jupyter_output><empty_output><jupyter_text>run a forward simulation with primary-secondary source<jupyter_code>src_inds2 = (
    (mesh.gridEy[:, 0] >= 0) & (mesh.gridEy[:, 0] <= radii.max()+eps) & 
    (mesh.gridEy[:, 2] >= -model.casing_l) & (mesh.gridEy[:, 2] <= 0)
)
src_inds2 = np.ones(mesh.nE, dtype=bool)

src_map2 = SourceMap(mesh, np.where(src_inds2)[0], sim.nT+1)
sigma = model.sigma_air * np.ones(mesh.nC)
sigma[mesh.gridCC[:, 2] < 0] = model.sigma_back
mysim_sm = MySimulation(
    mesh=mesh, time_steps=sim3D.time_steps, 
    solver=Solver, solver_opts=solver_opts, sigma=sigma, #sigma=model.sigma(mesh),
    source_map=src_map2, j_inds=j_inds
)
jd_x = (jd_3D[:mesh3D.nFx, :]).reshape(np.hstack([mesh3D.vnFx, np.r_[sim.nT+1]]), order="F")
jd_z = (jd_3D[mesh3D.vnF[:2].sum():, :]).reshape(np.hstack([mesh3D.vnFz, np.r_[sim.nT+1]]), order="F")

# grab a slice through theta
theta_ind = 3
jd = np.hstack([
    discretize.utils.mkvc(jd_x[:, theta_ind, :, 0]),
    discretize.utils.mkvc(jd_z[:, theta_ind, :, 0]),
])
jd.shape
src_inds3D = (
#     (mesh3D.gridEy[:, 0] >= 0) & (mesh3D.gridEy[:, 0] <= radii.max()+eps) &
    (mesh3D.gridEy[:, 1] == mesh3D.vectorCCy[3]) 
#     (mesh3D.gridEy[:, 2] >= -model.casing_l) & (mesh3D.gridEy[:, 2] <= 0)
)
smy3D = discretize.utils.mkvc(SMY[src_inds3D, :])
j_sm = mysim_sm.fields(smy3D, j0=discretize.utils.mkvc(jd))
tind = 40
fig, ax = plt.subplots(1, 2, figsize=(12, 4))

xlim = 1000*np.r_[-1, 1]
zlim = np.r_[-1500, 100]
sample_grid = np.r_[xlim.max()/50, 5.]
clim = np.r_[1e-12, 1e-6]

out = mesh.plotImage(
    mesh.aveF2CCV * j_sm[:, tind], 
    view="vec",
    vType="CCv",
    ax=ax[0], mirror=True,
    range_x=xlim, 
    range_y=zlim,
    sample_grid = sample_grid,
    pcolorOpts={"norm":LogNorm()},
    clim=clim, 
    stream_threshold=clim.min()
)

plt.colorbar(out[0], ax=ax[0])

ax[0].set_title(f"estimated, t={sim.times[tind]*1e3:1.1e}ms")

out2 = viewer.plot_cross_section(
    ax=ax[1], clim=clim, zlim=zlim,
    xlim=xlim,
    view='j', theta_ind=3, time_ind=tind, 
    model_key='approx_casing', show_cb=True, casing_outline=False,
    prim_sec="secondary"
#     stream_opts={"density":0.75, "color": "k", "arrowsize": 2}
)
# ax[1].set_ylim(np.r_[-max_depth, top])
ax[1].set_ylabel('z (m)')
ax[1].set_title(f"true , t={sim.times[tind]*1e3:1.1e}ms")

# for a in ax:
#     a.set_aspect(1)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.17501456027955736, 'fraction_numerical': 0.027872535152674932, 'mean_word_length': 3.2255229390050975, 'pattern_counts': {'":': 12, '<': 37, '<?xml version=': 0, '>': 38, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '99277', 'license': 'MIT', 'n_tokens_mistral': 10695, 'n_tokens_neox': 10342, 'n_words': 1730, 'path': 'notebooks/synthetic_source_inversion_SimPEG-electric-source.ipynb', 'repo_name': 'lheagy/heagy-2019-casing-ml', 'size': '24038'}