MemoryError using mesh.plot_slice() with a TreeMesh of around 500,000 cells

[aguspesce]

Hello,

I'm not sure if it's a bug or I'm doing something wrong.
I'm making a mesh with around 500,000 cells using TreeMesh and when I try to plot it, it returns a MemoryError.

I am using Discretize 0.10.0 and python 3.8.
This is my code to create the mesh:

import discretize
import SimPEG
import numpy as np
import matplotlib.pyplot as plt

casing_r = 0.2   # 10 cm of diameter  
casing_l = 300  #  
casing_t = 2/100  # 2cm thickness
casing_a = casing_r - casing_t / 2.0  # inner radius
casing_b = casing_r + casing_t / 2.0  # outer radius
casing_z = np.r_[-casing_l, 0.0]

electrode_a_location = np.array([0., 0., 0.])
electrode_b_location = np.array([5000., 0., 0.])

z_receivers = np.linspace(-300, -380, 81) 
x_receivers = np.zeros_like(z_receivers)
y_receivers = np.zeros_like(z_receivers)
receiver_locations = np.column_stack((x_receivers, y_receivers, z_receivers))

# Create the mesh
padding = electrode_b_location[0]*4
dh = casing_b/3
Lx = padding
Ly = padding 
Lz = padding
nbcx = 2 ** int(np.round(np.log(Lx / dh) / np.log(2.0)))
nbcy= 2 ** int(np.round(np.log(Ly / dh) / np.log(2.0)))
nbcz = 2 ** int(np.round(np.log(Lz / dh) / np.log(2.0)))
hx = [(dh, nbcx)]
hy = [(dh, nbcy)]
hz = [(dh, nbcz)]
mesh = discretize.TreeMesh([hx, hy, hz], origin="CCC")

mesh.refine_points(electrode_b_location, level=-3, padding_cells_by_level=[3, 2], finalize=False)
mesh.refine_points(electrode_a_location, level=-3, padding_cells_by_level=[3, 2], finalize=False)
mesh.refine_box([-casing_b, -casing_b, -casing_l], [casing_b, casing_b, 0], -1 , finalize=False)
mesh.refine_points(receiver_locations, level=-2, padding_cells_by_level=[3, 2], finalize=False)
mesh.finalize()

The resulting mesh is:

OcTreeMesh: 0.00% filled

Level : Number of cells               Mesh Extent               Cell Widths    
-----------------------           min     ,     max            min   ,   max   
  2   :       24             ---------------------------   --------------------
  3   :       272         x: -9175.040000000005,9175.039999966124      0.07   , 4587.520000002138
  4   :       332         y: -9175.040000000005,9175.039999966124      0.07   , 4587.520000002138
  5   :       364         z: -9175.040000000005,9175.039999966124   0.06999999999999318, 4587.520000002138
  6   :       352      
  7   :       436      
  8   :       496      
  9   :       724      
 10   :      1060      
 11   :      1960      
 12   :      3464      
 13   :      6540      
 14   :      12604     
 15   :      24732     
 16   :      53992     
 17   :      99264     
 18   :     274432     
-----------------------
Total :     481048  

When I try to plot it:

 _, ax = plt.subplots(figsize=(6, 4))
 mesh.plot_slice(np.ones(mesh.n_cells), normal='y', ax=ax, grid=False)
plt.show()

---------------------------------------------------------------------------
MemoryError                               Traceback (most recent call last)
Cell In[13], line 2
      1 _, ax = plt.subplots(figsize=(6, 4))
----> 2 mesh.plot_slice(mesh.cell_volumes, normal='y', ax=ax, grid=True)
      4 #ax.set_xlim(-1, 1)
      5 #ax.set_ylim(-310, 0)
      6 plt.show()

File ~/.mambaforge/envs/simpeg-env/lib/python3.8/site-packages/discretize/mixins/mpl_mod.py:682, in InterfaceMPL.plot_slice(self, v, v_type, normal, ind, slice_loc, grid, view, ax, clim, show_it, pcolor_opts, stream_opts, grid_opts, range_x, range_y, sample_grid, stream_threshold, stream_thickness, **kwargs)
    679     pcolor_opts["vmin"] = clim[0]
    680     pcolor_opts["vmax"] = clim[1]
--> 682 out = plotter(
    683     v,
    684     v_type=v_type,
    685     normal=normal,
    686     ind=ind,
    687     grid=grid,
    688     view=view,
    689     ax=ax,
    690     pcolor_opts=pcolor_opts,
    691     stream_opts=stream_opts,
    692     grid_opts=grid_opts,
    693     range_x=range_x,
    694     range_y=range_y,
    695     sample_grid=sample_grid,
    696     stream_threshold=stream_threshold,
    697     stream_thickness=stream_thickness,
    698     **kwargs,
    699 )
    700 if show_it:
    701     plt.show()

File ~/.mambaforge/envs/simpeg-env/lib/python3.8/site-packages/discretize/mixins/mpl_mod.py:2156, in InterfaceMPL.__plot_slice_tree(self, v, v_type, normal, ind, grid, view, ax, pcolor_opts, stream_opts, grid_opts, range_x, range_y, quiver_opts, **kwargs)
   2153 level_diff = self.max_level - temp_mesh.max_level
   2155 XS = [None, None, None]
-> 2156 XS[antiNormalInd[0]], XS[antiNormalInd[1]] = np.meshgrid(
   2157     cc_tensor[antiNormalInd[0]], cc_tensor[antiNormalInd[1]]
   2158 )
   2159 XS[normalInd] = np.ones_like(XS[antiNormalInd[0]]) * slice_loc
   2160 loc_grid = np.c_[XS[0].reshape(-1), XS[1].reshape(-1), XS[2].reshape(-1)]

File <__array_function__ internals>:200, in meshgrid(*args, **kwargs)

File ~/.mambaforge/envs/simpeg-env/lib/python3.8/site-packages/numpy/lib/function_base.py:5045, in meshgrid(copy, sparse, indexing, *xi)
   5042     output = np.broadcast_arrays(*output, subok=True)
   5044 if copy:
-> 5045     output = [x.copy() for x in output]
   5047 return output

File ~/.mambaforge/envs/simpeg-env/lib/python3.8/site-packages/numpy/lib/function_base.py:5045, in <listcomp>(.0)
   5042     output = np.broadcast_arrays(*output, subok=True)
   5044 if copy:
-> 5045     output = [x.copy() for x in output]
   5047 return output

MemoryError: Unable to allocate 512. GiB for an array with shape (262144, 262144) and data type float64

If I try to plot a slice of the mesh adding range_x=[-5, 5], range_y=[ -300, 0]) into the mesh.plot_slice() function, It also returns the same MemoryError.

Let me know how I can fix it.
Thank you!

[santisoler]

Thanks for reporting this bug, @aguspesce!

The issue is clearly coming from the plotting method trying to build a meshgrid out of coordinates of the cell centers of the underlying tensormesh. Since your smallest cells are too small in comparison to the range of the mesh, then the memory needed to allocate such meshgrid arrays is huge. I think we need to reimplement some of the pieces of this methods, since it shouldn't be required to allocate GBs of memory just to plot one slice of the mesh.

[jcapriot]

Oh boy, that is using a lot of memory to do this!
I've wanted to put more efficient intersection methods into the TreeMesh and/or give the user access to them. It already does efficient intersection tests for some simple geometric shapes in the refine functions, would just need to separate them into different functions and make the intersection tests rely on those.