#!/usr/bin/env python
# coding: utf-8

# # Grids: Non-Uniform Grids

# Some data cannot be easily represented on a grid of uniformly spaced vertices. It is still possible to create a grid object to represent such a dataset. 

# In[ ]:


get_ipython().run_line_magic('matplotlib', 'inline')

import astropy.units as u
import matplotlib.pyplot as plt
import numpy as np

from plasmapy.plasma import grids


# In[ ]:


grid = grids.NonUniformCartesianGrid(
    np.array([-1, -1, -1]) * u.cm, np.array([1, 1, 1]) * u.cm, num=(50, 50, 50)
)


# Currently, all non-uniform data is stored as an unordered 1D array of points. Therefore, although the dataset created above falls approximately on a Cartesian grid, its treatment is identical to a completely unordered set of points

# In[ ]:


grid.shape


# Many of the properties defined for uniform grids are inaccessible for non-uniform grids. For example, it is not possible to pull out an axis. However, the following properties still apply

# In[ ]:


print(f"Grid points: {grid.grid.shape}")
print(f"Units: {grid.units}")


# Properties can be added in the same way as on uniform grids.

# In[ ]:


Bx = np.random.rand(*grid.shape) * u.T
grid.add_quantities(B_x=Bx)
print(grid)


# ## Methods

# Many of the methods defined for uniform grids also work for non-uniform grids, however there is usually a substantial performance penalty in the non-uniform case.
# 
# For example, `grid.on_grid` behaves similarly. In this case, the boundaries of the grid are defined by the furthest point away from the origin in each direction.

# In[ ]:


pos = np.array([[0.1, -0.3, 0], [3, 0, 0]]) * u.cm
print(grid.on_grid(pos))


# The same definition is used to define the grid boundaries in `grid.vector_intersects` 

# In[ ]:


pt0 = np.array([3, 0, 0]) * u.cm
pt1 = np.array([-3, 0, 0]) * u.cm
pt2 = np.array([3, 10, 0]) * u.cm

print(f"Line from pt0 to pt1 intersects: {grid.vector_intersects(pt0, pt1)}")
print(f"Line from pt0 to pt2 intersects: {grid.vector_intersects(pt0, pt2)}")


# ## Interpolating Quantities

# Nearest-neighbor interpolation also works identically. However, volume-weighted interpolation is not implemented for non-uniform grids.

# In[ ]:


pos = np.array([[0.1, -0.3, 0], [0.5, 0.25, 0.8]]) * u.cm
print(f"Pos shape: {pos.shape}")
print(f"Position 1: {pos[0,:]}")
print(f"Position 2: {pos[1,:]}")

Bx_vals = grid.nearest_neighbor_interpolator(pos, "B_x")
print(f"Bx at position 1: {Bx_vals[0]:.2f}")


# In[ ]: