# coding: utf-8

# In[1]:

import numpy as np

from yt.fields.derived_field import     ValidateGridType,     ValidateParameter,     ValidateSpatial,     NeedsParameter
    
import yt


# In[2]:

from yt.funcs import     just_one


# In[3]:

sl_left = slice(None, -2, None)
sl_right = slice(2, None, None)
div_fac = 2.0

sl_center = slice(1, -1, None) 
ftype='gas'

#sl_left = slice(None,-2,None)

#sl_right = slice(1,-1,None)

#div_fac = 1.0


# In[4]:

def _Omega(field,data):
        if data.has_field_parameter("bulk_velocity"):
                bv = data.get_field_parameter("bulk_velocity").in_units("cm/s")
        else:
                bv = data.ds.arr(np.zeros(3), "cm/s")
        xv = data["gas","velocity_x"] - bv[0]
        yv = data["gas","velocity_y"] - bv[1]
        center = data.get_field_parameter('center')
        x_hat = data["x"] - center[0]
        y_hat = data["y"] - center[1]
        r = np.sqrt(x_hat*x_hat+y_hat*y_hat)
        x_hat /= r**2
        y_hat /= r**2

        return np.sqrt((xv*y_hat-yv*x_hat)**2)


yt.add_field("Omega", function=_Omega,take_log=False, units=r"1/yr",validators=[ValidateParameter('bulk_velocity')])


# In[5]:

def _dOmega_dx(field, data):
    
    f  = (data[ftype, "Omega"][sl_right,sl_center,sl_center] - data[ftype, "Omega"][sl_left,sl_center,sl_center])           / (div_fac*just_one(data["index", "dx"]))
    new_field = data.ds.arr(np.zeros_like(data[ftype, "velocity_z"],
                                          dtype=np.float64),
                            f.units)
    new_field[sl_center,sl_center,sl_center] = f
    return new_field
vort_validators = [ValidateSpatial(1,
                        [(ftype, "velocity_x"),
                         (ftype, "velocity_y"),
                         (ftype, "velocity_z")])]

yt.add_field(("gas","dOmega_dx"),
                   function=_dOmega_dx,
                   units="1/(pc*yr)",
                   validators=vort_validators)


# In[6]:

ds = yt.load('P010000/G-0000') 
ds.periodicity = (True, True, True)


# In[7]:

proj = yt.ProjectionPlot(ds, 'z', "density", weight_field="density")


# In[8]:

proj.zoom(5)


# In[9]:

proj1 = yt.SlicePlot(ds, 'z', "Omega")


# In[10]:

proj1.zoom(5)


# In[11]:

ds.add_gradient_fields(("gas","density"))
ds.add_gradient_fields(("gas","Omega"))


# In[12]:

proj3 = yt.SlicePlot(ds, 'z', "density_gradient_x")


# In[13]:

proj3.zoom(5)


# In[14]:

proj4 = yt.SlicePlot(ds, 'z', "Omega_gradient_x")


# In[15]:

proj4.zoom(5)


# In[16]:

proj5 = yt.SlicePlot(ds, 'z', "dOmega_dx")
proj5.zoom(5)


# In[ ]:




# In[ ]: