# THIRD PARTY
import agama
import astropy.coordinates as coord
import astropy.units as u
import numpy as np
from astropy.visualization import quantity_support
from galpy import df as gdf
from galpy import potential as gpot
# PROJECT-SPECIFIC
from discO import (
GaussianMeasurementErrorSampler,
PotentialFitter,
PotentialSampler,
conf,
)
from discO.plugin.agama.fitter import (
AGAMAMultipolePotentialFitter,
AGAMAPotentialFitter,
)
from discO.core.pipeline import Pipeline
from discO.core.residual import GridResidual
quantity_support();
mass = 1e12 * u.solMass
r0 = 10 * u.kpc # scale factor
agama.setUnits(mass=1, length=1, velocity=1)
hernquist_pot = gpot.HernquistPotential(amp=2 * mass, a=r0)
sampler = PotentialSampler(gdf.isotropicHernquistdf(hernquist_pot))
measurer = GaussianMeasurementErrorSampler(c_err=1*u.percent) # 10% error
fitter = PotentialFitter(None, key="agama", pot_type="multipole", symmetry="s", gridsizeR=50, lmax=0, mmax=0)
_x = np.linspace(-5, 5, num=50)
x, y, z = np.meshgrid(_x, _x, _x) * u.kpc
points = coord.CartesianRepresentation(x, y, z)
residualer = GridResidual(points, original_pot=hernquist_pot)
pipe = Pipeline(sampler=sampler, fitter=fitter, residualer=residualer)
pipe
resid = pipe(n=10000, frame=coord.Galactocentric(), random=0)
from mpl_toolkits.mplot3d import axes3d
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np
fig = plt.figure()
ax = fig.gca(projection='3d')
rar = np.abs(resid.represent_as(coord.SphericalRepresentation).vf_distance.flatten().value)
norm = mpl.colors.Normalize(vmin=min(rar), vmax=max(rar), clip=True)
mapper = cm.ScalarMappable(norm=norm, cmap=cm.inferno_r)
uvw = resid.vf_xyz.value * 1e13
qv = ax.quiver(
x, y, z, *uvw,
color=mapper.to_rgba(rar, alpha=0.1),
length=0.1, )
ax.set_xlim(-8, 8)
ax.set_ylim(-8, 8)
ax.set_zlim(-8, 8)
plt.colorbar(qv)
plt.show()