Extract core mesh from TreeMesh

[rowanc1]

The plotting for an octree is insanely slow. There is also some times that you want to work with a core mesh as a tensor mesh. This requires you to extract pieces and use those.

A quick sketch of the algorithm in 2D is something like:

# find the first cell of the tree mesh
x0_tensor = MT.gridCC[0, :]
for c in M:
    if np.allclose(c.center, x0_tensor):
        first_cell = c
        break
print(first_cell.center)
print(first_cell._pointer)
ptr = first_cell._pointer

index = np.zeros(MT.vnC, dtype=int)

for ii, nx in enumerate(range(MT.nCx)):
    for jj, ny in enumerate(range(MT.nCy)):
        # if you are not in the last level of the tree mesh 
        # you will have to skip up by that number
        # The last entry of the pointer refers to your level
        ptr_ij = [ptr[0] + nx, ptr[1] + ny, ptr[2]]
        index[ii, jj] = M._cc2i[M._index(ptr_ij)]

index = discretize.utils.mkvc(index)

And then you can plot it:

fig, axs = plt.subplots(1, 2, figsize=(15,8))
v = np.sin(M.gridCC[:, 0] * 2.0 * np.pi) * np.sin(M.gridCC[:, 1] * 2.0 * np.pi) + M.gridCC[:, 1]
M.plotImage(v, ax=axs[0])
axs[0].plot(MT.gridCC[:,0], MT.gridCC[:,1], 'kx')

MT.plotImage(v[index], ax=axs[1])

Would need a bunch of error checking, and extensions to 3D to make this robust.

cc @micmitch

[rowanc1]

A gotcha would be the pointer code if you are not on the lowest level. Should error in the M._cc2i lookup if the cell does not exist. you will have to step by the gap in ptr_ij = [ptr[0] + nx, ptr[1] + ny, ptr[2]] line You could change the enumerate(range(MT.nCy*level_from_bottom)[::level_from_bottom]) as a quick hack. i.e. level from bottom in this case is 1

[domfournier]

I am sure @jcapriot will have a better trick, but this would avoid looping and using built-in C functions:

Get info about the code region

hc = [h[0] for h in mesh.h]
core_cells = mesh.cell_volumes == mesh.cell_volumes.min()
min_locs = np.min(mesh.cell_centers[core_cells, :], axis=0)
max_locs = np.max(mesh.cell_centers[core_cells, :], axis=0)

Create an equivalent Tensor Mesh

nc = ((max_locs - min_locs) / hc + 1).astype(int)  # Could always add option to bufffer out
tensor = TensorMesh(
    [np.ones(c)*h for c, h in zip(nc, hc)], 
    x0=min_locs - hc/2
)

Get indices of treemesh

ind = mesh._get_containing_cell_indexes(tensor.cell_centers)

Plot model using tensor mesh

[thast]

Here is my own take at this, I interpolate to the base tensor mesh and use the extract_core_mesh utils.

from discretize import TensorMesh, TreeMesh
from discretize import utils as dsutils
import numpy as np

# original tree mesh
mesh = TreeMesh.read_UBC('./mesh.msh')
recovered_model = np.load('./recovered_model.npy')

#define core area
xyzlim = np.c_[
    [-200,+200],
    [-200,+200],
    [-100,+100],
].T

# create base tensormesh
tensormesh = TensorMesh(mesh.h)
tensormesh.x0 = mesh.x0

#extract core mesh
actcore, meshcore = dsutils.extract_core_mesh(xyzlim,tensormesh)

# interpolate tree mesh model to tensormesh
Fmodel = NearestNDInterpolator(mesh.gridCC,recovered_model)
modelcore = Fmodel(meshcore.gridCC)