def makePotentialModel(params):
# combined 4 components and the CMC represented by a single triaxial CylSpline potential
mmax = 12 # order of azimuthal Fourier expansion (higher order means better accuracy,
# but values greater than 12 *significantly* slow down the computation!)
pot_bary = agama.Potential(type='CylSpline',
density=agama.Density(
makeDensityModel(params),
makeCMC(0.2e10, 0.25, 0.05, 0.5)),
symmetry='t',
mmax=mmax,
gridsizeR=25,
gridsizez=25,
Rmin=0.1,
Rmax=40,
zmin=0.05,
zmax=20)
# flattened axisymmetric dark halo with the Einasto profile
pot_dark = agama.Potential(type='Multipole',
density='Spheroid',
axisratioz=0.8,
gamma=0,
beta=0,
outerCutoffRadius=1.84,
cutoffStrength=0.74,
densityNorm=0.0263e10,
gridsizer=26,
rmin=0.01,
rmax=1000,
lmax=8)
return agama.Potential(pot_bary, pot_dark)
def printoutInfo(model, iteration):
densDisk = model.components[0].getDensity()
densBulge = model.components[1].getDensity()
densHalo = model.components[2].getDensity()
pt0 = (2.0, 0, 0)
pt1 = (2.0, 0, 0.25)
pt2 = (0.0, 0, 2.0)
print("Disk total mass=%g, rho(R=2,z=0)=%g, rho(R=2,z=0.25)=%g" % \
(densDisk.totalMass(), densDisk.density(pt0), densDisk.density(pt1)))
print("Bulge total mass=%g, rho(R=0.5,z=0)=%g" % \
(densBulge.totalMass(), densBulge.density(0.4, 0, 0)))
print("Halo total mass=%g, rho(R=2,z=0)=%g, rho(R=0,z=2)=%g" % \
(densHalo.totalMass(), densHalo.density(pt0), densHalo.density(pt2)))
print("Potential at origin=-(%g)^2, total mass=%g" % \
((-model.potential.potential(0,0,0))**0.5, model.potential.totalMass()))
densDisk.export("dens_disk_iter" + str(iteration))
densBulge.export("dens_bulge_iter" + str(iteration))
densHalo.export("dens_halo_iter" + str(iteration))
model.potential.export("potential_iter" + str(iteration))
writeRotationCurve(
"rotcurve_iter" + str(iteration),
(
model.potential[1], # potential of the disk
agama.Potential(type='Multipole', lmax=6,
density=densBulge), # -"- bulge
agama.Potential(type='Multipole', lmax=6,
density=densHalo))) # -"- halo
def plotVcirc(model, iteration):
rhos = (model.components[0].getDensity(), model.components[1].getDensity())
pots = ( # recompute potentials of both components separately, using a multipole Poisson solver
agama.Potential(type='Multipole',
lmax=6,
density=rhos[0],
rmin=1e-4,
rmax=0.1),
agama.Potential(type='Multipole',
lmax=12,
density=rhos[1],
rmin=1e-3,
rmax=1.0))
r = numpy.linspace(0.0, 1.0, 1001)
xyz = numpy.column_stack((r, r * 0, r * 0))
vcomp2 = numpy.column_stack([-pot.force(xyz)[:, 0] * r for pot in pots])
vtot2 = numpy.sum(vcomp2, axis=1)
#print('NSC total mass=%g'%rhos[0].totalMass())
print('NSD total mass=%g' % rhos[1].totalMass())
# plot the circular-velocity curves for both NSC and NSD separately, and their sum
plt.figure(figsize=(6, 4))
ax = plt.axes([0.15, 0.15, 0.8, 0.8])
ax.plot(r, vtot2**0.5, color='k', label='total')
ax.plot(r, vcomp2[:, 0]**0.5, '--', color='k', label='NSC')
ax.plot(r, vcomp2[:, 1]**0.5, ':', color='k', label='NSD')
ax.set_xlim(0, 1)
ax.set_ylim(0, 140)
ax.set_xlabel(r'$R\, {\rm [kpc]}$', fontsize=16)
ax.set_ylabel(r'$v_{\rm circ}\, {\rm [km/s]}$', fontsize=16)
ax.grid()
ax.legend()
plt.savefig('nsd_vcirc_iter%i.pdf' % iteration)
plt.close()
def printoutInfo(model, iteration):
densDisk = model.components[0].getDensity()
densBulge= model.components[1].getDensity()
densHalo = model.components[2].getDensity()
pt0 = (2.0, 0, 0)
pt1 = (2.0, 0, 0.25)
pt2 = (0.0, 0, 2.0)
print("Disk total mass=%g, rho(R=2,z=0)=%g, rho(R=2,z=0.25)=%g" % \
(densDisk.totalMass(), densDisk.density(pt0), densDisk.density(pt1)))
print("Bulge total mass=%g, rho(R=0.5,z=0)=%g" % \
(densBulge.totalMass(), densBulge.density(0.4, 0, 0)))
print("Halo total mass=%g, rho(R=2,z=0)=%g, rho(R=0,z=2)=%g" % \
(densHalo.totalMass(), densHalo.density(pt0), densHalo.density(pt2)))
# report only the potential of stars+halo, excluding the potential of the central BH (0th component)
pot0 = model.potential.potential(0,0,0) - model.potential[0].potential(0,0,0)
print("Potential at origin=-(%g)^2, total mass=%g" % ((-pot0)**0.5, model.potential.totalMass()))
densDisk. export("dens_disk_" +iteration)
densBulge.export("dens_bulge_"+iteration)
densHalo. export("dens_halo_" +iteration)
model.potential.export("potential_"+iteration)
# separate the contributions of bulge and halo, which are normally combined
# into the Multipole potential of all spheroidal components
writeRotationCurve("rotcurve_"+iteration, (
model.potential[0], # potential of the BH
model.potential[2], # potential of the disk
agama.Potential(type='Multipole', lmax=6, density=densBulge), # -"- bulge
agama.Potential(type='Multipole', lmax=6, density=densHalo) ), # -"- halo
('BH', 'Disk', 'Bulge', 'Halo') )
def _gen_axisymmetric_(self):
import agama
potential_cache_file = self.cache_directory + '/potential_id'+str(self.snap_index)
potential_cache_file += '_' + self.sim_name + '_pot'
try:
self.potential = agama.Potential(file=potential_cache_file)
except:
star_position = self.first_snapshot['star'].prop('host.distance.principal')
gas_position = self.first_snapshot['gas'].prop('host.distance.principal')
dark_position = self.first_snapshot['dark'].prop('host.distance.principal')
star_mass = self.first_snapshot['star']['mass']
gas_mass = self.first_snapshot['gas']['mass']
dark_mass = self.first_snapshot['dark']['mass']
position = np.concatenate((star_position, gas_position))
mass = np.concatenate((star_mass, gas_mass))
#TODO make these user-controllable
self.pdark = agama.Potential(type="Multipole",
particles=(dark_position, dark_mass),
symmetry='a', gridsizeR=20, lmax=2)
self.pbar = agama.Potential(type="CylSpline",
particles=(position, mass),
symmetry='a', gridsizer=20, gridsizez=20,
mmax=0, Rmin=0.2,
Rmax=50, Zmin=0.02, Zmax=10)
self.potential = agama.Potential(self.pdark, self.pbar)
self.potential.export(potential_cache_file)
return None
def __init__(self, filename):
'''
Initialize starting values of scaled parameters and and define upper/lower limits on them;
also obtain the true values by analyzing the input file name
'''
self.initValues = numpy.array(
[9., 0., 0.5, 4.0, 1.0, 6.0, 1.5, 1.0, 1.0, 1.0, 1.])
self.minValues = numpy.array(
[6., -1., 0.0, 2.2, .25, 3.2, 0.0, 0.5, 0.1, 0.1, -1.])
self.maxValues = numpy.array(
[10., 2., 1.9, 6.0, 2.5, 12., 2.8, 2.0, 2.8, 2.8, 3.])
self.labels = ( \
r'$\log_{10}(\rho_1)$', r'$\log_{10}(R_{scale})$', r'$\gamma$', r'$\beta$', r'$\alpha$', \
r'$B_{DF}$', r'$\Gamma_{DF}$', r'$\eta$', r'$g_r$', r'$h_r$', r'$\log_{10}(J_0)$')
self.numPotParams = 5 # potential params come first in the list
self.numDFParams = 6 # DF params come last in the list
self.scaleRadius = 1. # fixed radius r_1 at which the DM density is constrained (rho_1)
# parse true values
m = re.match(
r'gs(\d+)_bs(\d+)_rcrs([a-z\d]+)_rarc([a-z\d]+)_([a-z]+)_(\d+)mpc3',
filename)
n = re.match(r'data_([ch])_rh(\d+)_rs(\d+)_gs(\d+)', filename)
if m:
self.truePotential = agama.Potential(
type='SpheroidDensity',
densityNorm=float(m.group(6)) * 1e6,
scaleRadius=1.0,
gamma=1.0 if m.group(5) == 'cusp' else 0.0,
beta=3.0)
self.scaleRadius = float(m.group(3)) * 0.01
self.tracerParams = dict(type='SpheroidDensity',
densityNorm=1.0,
scaleRadius=self.scaleRadius,
gamma=float(m.group(1)) * 0.01,
beta=float(m.group(2)) * 0.1)
# normalize the tracer density profile to have a total mass of unity
self.tracerParams["densityNorm"] /= agama.Density(
**self.tracerParams).totalMass()
self.tracerDensity = agama.Density(**self.tracerParams)
elif n:
self.truePotential = agama.Potential(
type='SpheroidDensity',
densityNorm=3.021516e7 if n.group(1) == 'c' else 2.387329e7,
scaleRadius=float(n.group(2)),
gamma=0.0 if n.group(1) == 'c' else 1.0,
beta=4.0)
self.scaleRadius = float(n.group(3)) * 0.01
self.tracerParams = dict(type='SpheroidDensity',
densityNorm=1.0,
scaleRadius=self.scaleRadius,
gamma=float(m.group(4)) * 0.1,
beta=5.0)
self.tracerParams["densityNorm"] /= agama.Density(
**self.tracerParams).totalMass()
self.tracerDensity = agama.Density(**self.tracerParams)
else:
print "Can't determine true parameters!"
def construct_potential_of_type(snap, ptype_list, center=np.array([0, 0, 0]), fout_base=None):
if isinstance(ptype_list, str):
ptype_list = (ptype_list,)
pot_list = []
for ptype in ptype_list:
part = getattr(snap, 'PartType'+str(ptype))
if len(part.data) == 0: # contains no particles
continue
pos = part['Coordinates']
if
if ptype == 'PartType1' or ptype == 'PartType3':
pot = agama.Potential(type='Multipole', particles=(pos, mass),
symmetry='a', gridsizeR=20, lmax=2)
else:
pot = agama.Potential(type='CylSpline', particles=(pos, mass),
symmetry='a', gridsizer=20, gridsizez=20, mmax=0,
Rmin=0.2, Rmax=50, Zmin=0.02, Zmax=10)
pot_list.append(pot)
return pot_list
def compute_potential(path, snapnum, name, center=np.array([0, 0, 0]), output_dir='data/'):
# try loading snapshot
sn = arepo.Snapshot(path+'output/', snapnum, combineFiles=True)
return sn
if __name__ == '__main__':
basepath = '../../runs/'
fid_g1 = 'fid-disp1.0-fg0.1'
fid_d15_g1 = 'fid-disp1.5-fg0.1'
# look to see if we are on my macbook or on the cluster
if sys.platform == 'darwin':
pair_list = [(fid_g1, 'lvl5')]
else:
pair_list = [(fid_g1, 'lvl5'), (fid_g1, 'lvl4'), (fid_g1, 'lvl3')]
name_list = [ p[0] + '-' + p[1] for p in pair_list]
path_list = [basepath + p[0] + '/' + p[1] for p in pair_list]
# nsnap_list = [len(glob.glob(path+'/output/snapdir*/*.0.hdf5')) for path in path_list]
snap_list = np.arange(0, 100, 10)
snapnum_list = [10, 50, 100, 150, 200, 300, 400, 500, 600]
for path, name in zip(tqdm(path_list), name_list):
for snapnum in snap_list:
out = compute_potential(path, snapnum, name, center=np.array([200, 200, 200]), output_dir='pot/')
def createModel(self, params):
'''
create a model (potential and DF) specified by the given [scaled] parameters
'''
potential = agama.Potential(type='Spheroid',
densityNorm=10**params[0],
scaleRadius=10**params[1],
gamma=params[2],
beta=params[3],
alpha=params[4])
density = agama.Density(type='Spheroid',
scaleRadius=10**params[self.numPotParams + 0],
gamma=params[self.numPotParams + 1],
beta=params[self.numPotParams + 2],
alpha=params[self.numPotParams + 3])
df = agama.DistributionFunction(type='QuasiSpherical',
potential=potential,
density=density,
beta0=params[self.numPotParams + 4],
r_a=10**params[self.numPotParams + 5])
# check if the DF is everywhere nonnegative
j = numpy.logspace(-5, 10, 200)
if any(df(numpy.column_stack((j, j * 0 + 1e-10, j * 0))) <= 0):
raise ValueError("Bad DF")
return potential, df
def setup_class(cls):
"""Setup fixtures for testing."""
agama.setUnits(mass=1, length=1, velocity=1) # FIXME! bad
cls.potential = agama.Potential(type="Plummer")
# set up the rest
super().setup_class()
def convertCoefsToAgamaPotential(r0,
Acos,
Asin=None,
filename='tmppotential.ini'):
'''
convert the arrays of cosine and sine coefs from galpy or gala into the agama input format,
and create the equivalent agama BasisSet potential
'''
nmax = Acos.shape[0] - 1
lmax = Acos.shape[1] - 1
assert Acos.shape[2] == lmax + 1
if Asin is None: Asin = numpy.zeros(Acos.shape)
with open(filename, 'w') as inifile:
inifile.write(
'[Potential]\ntype=BasisSet\nnmax=%i\nlmax=%i\nr0=%g\nCoefficients\n#Phi\n#(array)\n'
% (nmax, lmax, r0))
for n in range(nmax + 1):
inifile.write(str(n))
for l in range(lmax + 1):
# first come sine terms in reverse order: m=l, l-1, ..., 1
for m in range(l):
inifile.write('\t%.15g' % (Asin[n, l, l - m] * 0.5**0.5))
# then come cosine terms in normal order: m=0, 1, ..., l
for m in range(l + 1):
inifile.write('\t%.15g' % (Acos[n, l, m] *
(0.5**0.5 if m > 0 else 1)))
inifile.write('\n')
return agama.Potential(filename)
def __init__(self,*args,**kwargs):
"""
NAME:
__init__
PURPOSE:
initialize a potential from parameters provided in an INI file
or as named arguments to the constructor (see below).
INPUT:
normalize - if True, normalize such that vc(1.,0.)=1., or,
if given as a number, such that the force is this fraction of the force
necessary to make vc(1.,0.)=1.
HISTORY:
2014-12-05 EV
"""
Potential.__init__(self,amp=1.)
normalize=False
for key, value in kwargs.items():
if key=="normalize":
normalize=value
del kwargs[key]
self._pot = agama.Potential(*args,**kwargs)
if normalize or \
(isinstance(normalize,(int,float)) \
and not isinstance(normalize,bool)):
self.normalize(normalize)
self.hasC= False
self.hasC_dxdv=False
def createSpiralPotential(numberOfArms,
surfaceDensity,
scaleRadius,
scaleHeight,
pitchAngle,
phi0=0):
'''
Create a CylSpline approximation for a spiral arm potential from Cox&Gomez(2002).
The density generated by this potential approximately follows an exponential
disk profile with the amplitude of surface density variation
Sigma(R) = surfaceDensity * exp( - R / scaleRadius ),
and isothermal vertical profile
rho(R,z) = Sigma(R) / (4h) * sech^2( z / (2h) ),
where h is the scaleHeight (note the extra factor 1/2 w.r.t. the original paper,
which is introduced to match the definition of the Disk profile in Agama).
The density at the crest of the spiral arm is ~ 3 * rho(R,z=0),
while in the inter-arm region it is ~ (-1.5 to -1) * rho(R,z=0).
The spiral is logarithmic, so that the azimuthal angle (of one of the arms) is
phi(R) ~ phi0 + ln(R/scaleRadius) / tan(pitchAngle).
'''
def potfnc(xyz):
R = (xyz[:, 0]**2 + xyz[:, 1]**2)**0.5
z = xyz[:, 2]
phi = numpy.arctan2(xyz[:, 1], xyz[:, 0])
prefac = -4 * numpy.pi * agama.G * surfaceDensity * numpy.exp(
-R / scaleRadius)
gamma = numberOfArms * (
phi - numpy.log(R / scaleRadius) / numpy.tan(pitchAngle) - phi0)
Phi = numpy.zeros(len(R))
for n in range(1, 4):
K_n = n * numberOfArms / R / numpy.sin(pitchAngle)
K_n_H = K_n * 2 * scaleHeight
B_n = K_n_H * (1 + 0.4 * K_n_H)
D_n = 1 / (1 + 0.3 * K_n_H) + K_n_H
C_n = [8. / 3 / numpy.pi, 0.5, 8. / 15 / numpy.pi
][n - 1] # amplitudes of the three harmonic terms
Phi += prefac * (C_n / D_n / K_n * numpy.cos(n * gamma) *
(numpy.cosh(K_n * z / B_n))**-B_n)
return numpy.nan_to_num(
Phi) # VERY IMPORTANT is to make sure it never produces a NaN
# now create a CylSpline potential approximating this user-defined profile,
# paying special attention to the grid parameters
return agama.Potential(
type='CylSpline',
potential=potfnc,
Rmin=0.01 * scaleRadius, # (these could be adjusted if needed)
Rmax=10.0 * scaleRadius,
zmin=0.25 * scaleHeight,
zmax=10.0 * scaleHeight,
mmax=3 *
numberOfArms, # each arm is represented by three harmonic terms: m, m*2, m*3
symmetry='bisymmetric' if numberOfArms %
2 == 0 else 4, # 4 means the z-reflection symmetry
gridSizeZ=20,
# rule of thumb for the radial grid spacing is to resolve
# the change in the pitch angle of the spiral with one grid cell
gridSizeR=max(
25,
numpy.log(10.0 / 0.01) / numpy.tan(pitchAngle) * numberOfArms))
def createModel(self, params):
'''
create a model (potential and DF) specified by the given [scaled] parameters
'''
potential = agama.Potential(type='Spheroid',
densityNorm=10**params[0],
scaleRadius=10**params[1],
gamma=params[2],
beta=params[3],
alpha=params[4])
# first create an un-normalized DF
dfparams = dict(type='DoublePowerLaw',
slopeOut=params[self.numPotParams + 0],
slopeIn=params[self.numPotParams + 1],
steepness=params[self.numPotParams + 2],
coefJrOut=params[self.numPotParams + 3],
coefJzOut=(3 - params[self.numPotParams + 3]) / 2,
coefJrIn=params[self.numPotParams + 4],
coefJzIn=(3 - params[self.numPotParams + 4]) / 2,
j0=10**params[self.numPotParams + 5],
norm=1.)
# compute its total mass
totalMass = agama.DistributionFunction(**dfparams).totalMass()
# and now normalize the DF to have a unit total mass
dfparams["norm"] = 1. / totalMass
df = agama.DistributionFunction(**dfparams)
return potential, df
def setup_class(cls):
"""Setup fixtures for testing."""
super().setup_class()
# now agama stuff
# override super
agama.setUnits(mass=1, length=1, velocity=1) # FIXME! bad
cls.potential = agama.Potential(type="Plummer")
def main():
pot = agama.Potential(type="Dehnen", mass=1, scaleRadius=1.)
actf = agama.ActionFinder(pot)
particles, masses = createHernquistModel(100000)
actions = actf(particles)
# do a parameter search to find best-fit distribution function describing these particles
initparams = numpy.array([2.0, 4.0, 1.0, 1.0, 0.0])
result = minimize(model_search_fnc,
initparams,
args=(actions, ),
method='Nelder-Mead',
options=dict(maxiter=1000, maxfev=1000, disp=True))
# explore the parameter space around the best-fit values using the MCMC chain
try:
import matplotlib.pyplot as plt, emcee, corner
except ImportError as ex:
print ex, "\nYou need to install 'emcee' and 'corner' packages"
print 'Starting MCMC'
ndim = len(initparams)
nwalkers = 16 # number of parallel walkers in the chain
nsteps = 300 # number of steps in MCMC chain
nburnin = 100 # number of initial steps to discard
# initial coverage of parameter space - around the best-fit solution with a small dispersion
initwalkers = [
result.x + 0.01 * numpy.random.randn(ndim) for i in range(nwalkers)
]
sampler = emcee.EnsembleSampler(nwalkers,
ndim,
model_search_emcee,
args=(actions, ))
sampler.run_mcmc(initwalkers, nsteps)
# show the time evolution of parameters carried by the ensemble of walkers (time=number of MC steps)
fig, axes = plt.subplots(ndim + 1, 1, sharex=True)
for i in range(ndim):
axes[i].plot(sampler.chain[:, :, i].T, color='k', alpha=0.5)
axes[i].set_ylabel(labels[i])
# last panel shows the evolution of log-likelihood for the ensemble of walkers
axes[-1].plot(sampler.lnprobability.T, color='k', alpha=0.5)
axes[-1].set_ylabel('log(L)')
maxloglike = numpy.max(sampler.lnprobability)
axes[-1].set_ylim(maxloglike - 3 * ndim, maxloglike)
fig.tight_layout(h_pad=0.)
plt.show()
# show the posterior distribution of parameters
samples = sampler.chain[:, nburnin:, :].reshape((-1, ndim))
trueval = (1.56, 1.55, 5.29, 1.22, 1.56) # believed to be best-fit values
corner.corner(samples, \
labels=labels, quantiles=[0.16, 0.5, 0.84], truths=trueval)
plt.show()
print "Acceptance fraction: ", numpy.mean(
sampler.acceptance_fraction) # should be in the range 0.2-0.5
print "Autocorrelation time: ", sampler.acor # should be considerably shorter than the total number of steps
def printoutInfo(model, iteration):
densDisk = model.components[0].getDensity()
densBulge= model.components[1].getDensity()
densHalo = model.components[2].getDensity()
pt0 = (solarRadius, 0, 0)
pt1 = (solarRadius, 0, 1)
print("Disk total mass=%g Msun, rho(Rsolar,z=0)=%g, rho(Rsolar,z=1kpc)=%g Msun/pc^3" % \
(densDisk.totalMass(), densDisk.density(pt0)*1e-9, densDisk.density(pt1)*1e-9)) # per pc^3, not kpc^3
print("Halo total mass=%g Msun, rho(Rsolar,z=0)=%g, rho(Rsolar,z=1kpc)=%g Msun/pc^3" % \
(densHalo.totalMass(), densHalo.density(pt0)*1e-9, densHalo.density(pt1)*1e-9))
print("Potential at origin=-(%g km/s)^2, total mass=%g Msun" % \
((-model.potential.potential(0,0,0))**0.5, model.potential.totalMass()))
densDisk.export ("dens_disk_" +iteration);
densBulge.export("dens_bulge_"+iteration);
densHalo.export ("dens_halo_" +iteration);
model.potential.export("potential_"+iteration);
writeRotationCurve("rotcurve_"+iteration, (model.potential[1], # disk potential (CylSpline)
agama.Potential(type='Multipole', lmax=6, density=densBulge), # -"- bulge
agama.Potential(type='Multipole', lmax=6, density=densHalo) ) ) # -"- halo
def updatePotential(self):
print("Updating potential...")
# sort out density and potential components into several groups
densitySph = []
densityCyl = []
potentials = []
for component in self.components:
dens = component.getDensity()
if dens is not None:
if component.disklike: densityCyl.append(dens)
else: densitySph.append(dens)
else:
potentials.append(component.getPotential())
# create a single Multipole potential for all non-disk-like density components
if len(densitySph) > 0:
potentials.append(
agama.Potential(type='Multipole',
density=agama.Density(*densitySph),
gridsizer=self.sizeradialsph,
rmin=self.rminsph,
rmax=self.rmaxsph,
lmax=self.lmaxangularsph,
mmax=0,
symmetry='a'))
# create a single CylSpline potential representing all disk-like density components
if len(densityCyl) > 0:
potentials.append(
agama.Potential(type='CylSpline',
density=agama.Density(*densityCyl),
gridsizer=self.sizeradialcyl,
rmin=self.rmincyl,
rmax=self.rmaxcyl,
gridsizez=self.sizeverticalcyl,
zmin=self.zmincyl,
zmax=self.zmaxcyl,
mmax=0,
symmetry='a'))
# combine all potential components and reinitialize the action finder
self.potential = agama.Potential(*potentials)
print("Updating action finder...")
self.af = agama.ActionFinder(self.potential)
def test(Component, SelfConsistentModel):
# test a two-component spherical model (to speed up things, do not use disky components);
# the first component is an isotropic DF of a NFW halo with a cutoff,
# and the second one represents a more concentrated baryonic component,
# which will cause the adiabatic contraction of the halo
# when both components are iterated to achieve equilibrium
params_comp1 = dict(disklike=False,
rminSph=0.01,
rmaxSph=100.0,
sizeRadialSph=21,
lmaxAngularSph=0)
params_comp2 = dict(disklike=False,
rminSph=0.01,
rmaxSph=10.0,
sizeRadialSph=16,
lmaxAngularSph=0)
params_scm = dict(rminSph=0.005,
rmaxSph=200,
sizeRadialSph=25,
lmaxAngularSph=0)
potential_init = agama.Potential(type='spheroid',
gamma=1,
beta=3,
mass=20.0,
scaleRadius=5.0,
outerCutoffRadius=40.0)
df_comp1 = agama.DistributionFunction(type='quasispherical',
density=potential_init,
potential=potential_init)
df_comp2 = agama.DistributionFunction(type='doublepowerlaw',
norm=12.0,
J0=1.0,
coefJrIn=1.,
coefJzIn=1.,
coefJrOut=1.,
coefJzOut=1.,
slopeIn=1,
slopeOut=6)
model = SelfConsistentModel(**params_scm)
model.components.append(Component(df=df_comp1, **params_comp1))
model.components.append(Component(df=df_comp2, **params_comp2))
model.potential = potential_init
for it in range(5):
model.iterate()
#model.components[0].getDensity().export('comp1')
#model.components[1].getDensity().export('comp2')
return agama.GalaxyModel(model.potential,
df_comp1).moments([1, 0.5, 0.3])
def createPotential(self, params):
rho1 = 10**params[0]
r0 = 10**params[1]
gamma = params[2]
beta = params[3]
alpha = params[4]
r1r0 = self.scaleRadius / r0
rho0 = rho1 * (1 + r1r0**alpha)**((beta - gamma) / alpha) * r1r0**gamma
return agama.Potential(type='SpheroidDensity',
densityNorm=rho0,
scaleRadius=r0,
gamma=gamma,
beta=beta,
alpha=alpha)
def setup_class(cls):
"""Setup fixtures for testing."""
super().setup_class()
# make potential
cls.potential = agama.Potential(
type="Spheroid",
mass=1e12,
scaleRadius=10,
gamma=1,
alpha=1,
beta=4,
cutoffStrength=0,
)
cls.inst = cls.obj(AGAMAPotentialWrapper(cls.potential))
def _gen_axisymmetric_(self, all_snaps=False):
agama.setUnits(mass=1, length=1, velocity=1)
if all_snaps:
pot_file_list = []
for index in self.snapshot_indices:
pot_file_list.append(self._pot_cache_file_(index))
else:
pot_file_list = (self._pot_cache_file_(self.startnum), )
for i, file in enumerate(pot_file_list):
try:
self.pdark = agama.Potential(file=file + '_dark')
self.pbar = agama.Potential(file=file + '_bar')
self.potential = agama.Potential(self.pdark, self.pbar)
except:
star_position = self.snapshots[i]['star'].prop(
'host.distance.principal')
gas_position = self.snapshots[i]['gas'].prop(
'host.distance.principal')
dark_position = self.snapshots[i]['dark'].prop(
'host.distance.principal')
star_mass = self.snapshots[i]['star']['mass']
gas_mass = self.snapshots[i]['gas']['mass']
dark_mass = self.snapshots[i]['dark']['mass']
position = np.concatenate((star_position, gas_position))
mass = np.concatenate((star_mass, gas_mass))
#TODO make these user-controllable
self.pdark = agama.Potential(type="Multipole",
particles=(dark_position,
dark_mass),
symmetry='a',
gridsizeR=20,
lmax=2)
self.pbar = agama.Potential(type="CylSpline",
particles=(position, mass),
symmetry='a',
gridsizer=20,
gridsizez=20,
mmax=0,
Rmin=0.2,
Rmax=50,
Zmin=0.02,
Zmax=10)
self.pdark.export(file + '_dark')
self.pbar.export(file + '_bar')
self.potential = agama.Potential(self.pdark, self.pbar)
self.potential.export(file)
if self.axisymmetric_tevolve:
self._gen_agama_interpolator_()
return None
iniPotenDisk = dict(ini.items("Potential disk"))
iniPotenBH = dict(ini.items("Potential BH"))
iniDFDisk = dict(ini.items("DF disk"))
iniSCMHalo = dict(ini.items("SelfConsistentModel halo"))
iniSCMBulge = dict(ini.items("SelfConsistentModel bulge"))
iniSCMDisk = dict(ini.items("SelfConsistentModel disk"))
iniSCM = dict(ini.items("SelfConsistentModel"))
# initialize the SelfConsistentModel object (only the potential expansion parameters)
model = agama.SelfConsistentModel(**iniSCM)
# create initial density profiles of all components
densityDisk = agama.Density(**iniPotenDisk)
densityBulge = agama.Density(**iniPotenBulge)
densityHalo = agama.Density(**iniPotenHalo)
potentialBH = agama.Potential(**iniPotenBH)
# add components to SCM - at first, all of them are static density profiles
model.components.append(agama.Component(density=densityDisk, disklike=True))
model.components.append(agama.Component(density=densityBulge, disklike=False))
model.components.append(agama.Component(density=densityHalo, disklike=False))
model.components.append(agama.Component(potential=potentialBH))
# compute the initial potential
model.iterate()
printoutInfo(model,'init')
# construct the DF of the disk component, using the initial (non-spherical) potential
dfDisk = agama.DistributionFunction(potential=model.potential, **iniDFDisk)
# initialize the DFs of spheroidal components using the Eddington inversion formula
# for their respective density profiles in the initial potential
except:
print("Input snapshot files are not available; " \
"you may create them by running example_self_consistent_model.py")
exit()
print("%g s to load %d disk particles (total mass=%g Msun) " \
"and %d halo particles (total mass=%g Msun)" % \
( time.clock()-tbegin, \
diskParticles.shape[0], numpy.sum(diskParticles[:,6]), \
haloParticles.shape[0], numpy.sum(haloParticles[:,6]) ) )
#3. create an axisymmetric potential from these snapshots
try:
#3a. try to load potentials from previously stored text files instead of computing them
diskPot = agama.Potential("model_stars_final.pot")
haloPot = agama.Potential("model_dm_final.pot")
except:
# 3b: these files don't exist on the first run, so we have to create the potentials
tbegin = time.clock()
haloPot = agama.Potential( \
type="Multipole", particles=(haloParticles[:,0:3], haloParticles[:,6]), \
symmetry='a', gridsizeR=20, lmax=2)
print("%f s to init %s potential for the halo; value at origin=%f (km/s)^2" % \
((time.clock()-tbegin), haloPot.name(), haloPot.potential(0,0,0)))
tbegin = time.clock()
# manually specify the spatial grid for the disk potential,
# although one may rely on the automatic choice of these parameters (as we did for the halo)
diskPot = agama.Potential( \
type="CylSpline", particles=(diskParticles[:,0:3], diskParticles[:,6]), \
densityHalo = agama.Density(**iniPotenHalo)
# add components to SCM - at first, all of them are static density profiles
model.components.append(agama.Component(density=densityDisk,
disklike=True))
model.components.append(
agama.Component(density=densityBulge, disklike=False))
model.components.append(
agama.Component(density=densityHalo, disklike=False))
# compute the initial potential
model.iterate()
writeRotationCurve(
"rotcurve_init",
(model.potential[1],
agama.Potential(type='Multipole', lmax=0, density=densityBulge),
agama.Potential(type='Multipole', lmax=0, density=densityHalo)))
# construct the DF of the disk component, using the initial (non-spherical) potential
dfDisk = agama.DistributionFunction(potential=model.potential, **iniDFDisk)
# initialize the DFs of spheroidal components using the Eddington inversion formula
# for their respective density profiles in the initial potential
dfBulge = agama.DistributionFunction(type='QuasiSpherical',
potential=model.potential,
density=densityBulge)
dfHalo = agama.DistributionFunction(type='QuasiSpherical',
potential=model.potential,
density=densityHalo)
print("\033[1;33m**** STARTING ITERATIVE MODELLING ****\033[0m\nMasses (computed from DF): " \
"Mdisk=%g, Mbulge=%g, Mhalo=%g" % (dfDisk.totalMass(), dfBulge.totalMass(), dfHalo.totalMass()))
# here we use the N-body model generated by another example program: example_self_consistent_model3.py
# among other things, it outputs two N-body snapshot files - one for the disk, the other for the bulge component
try:
print('Reading input snapshot')
snapshot1 = agama.readSnapshot('model_disk_final')
snapshot2 = agama.readSnapshot('model_bulge_final')
posvel = numpy.vstack((snapshot1[0], snapshot2[0])) # 2d Nx6 array of positions and velocities
mass = numpy.hstack((snapshot1[1], snapshot2[1])) # 1d array of N particle masses
# if your N-body snapshot is contained in a single file, just load it and assign posvel,mass arrays as specified above
except:
print('You need to generate N-body snapshots by running example_self_consistent_model3.py')
exit()
# convert the N-body model (which was set up in N-body units with G=1) to observational units defined at the beginning of this script
rscale = 30.0 # [REQ] 1 length unit of the N-body snapshot corresponds to this many length units of this script (arcseconds)
mscale = 4e10 # [REQ] 1 mass unit of the snapshot corresponds to this many mass units of this script (solar masses)
G = -agama.Potential(type='Plummer').potential(0,0,0) # numerical value of G in current units
vscale = (G * mscale / rscale)**0.5 # same for the N-body velocity unit => km/s
posvel[:, 0:3] *= rscale
posvel[:, 3:6] *= vscale
mass *= mscale
# pre-step ([OPT] - can use only for axisymmetric systems): create several rotated copies of the input snapshot to reduce Poisson noise
nrot = 9 # [OPT] number of rotation angles
posvel_stack = []
print('Creating %d rotated copies of input snapshot' % nrot)
for ang in numpy.linspace(0, numpy.pi, nrot):
sina, cosa = numpy.sin(ang), numpy.cos(ang)
posvel_stack.append( numpy.column_stack((
posvel[:,0] * cosa + posvel[:,1] * sina,
posvel[:,1] * cosa - posvel[:,0] * sina,
posvel[:,2],
np.cos(theta[1])]])
R_z = np.array([[np.cos(theta[2]), -np.sin(theta[2]), 0],
[np.sin(theta[2]), np.cos(theta[2]), 0], [0, 0, 1]])
R = np.dot(R_x, np.dot(R_y, R_z))
# this is the real one
# R = np.dot(R_z, np.dot( R_y, R_x ))
return R
def euler_rotate(theta, vec):
mat = euler_matrix(theta)
return np.transpose(np.tensordot(mat, np.transpose(vec), axes=1))
potential = agama.Potential(file='fiducial_pot')
af = agama.ActionFinder(potential, interp=False)
# read in simulated positions and velocities
posfile = 'pos_' + sys.argv[1] + '.npy'
velfile = 'vel_' + sys.argv[1] + '.npy'
pos_full = np.load(posfile)
vel_full = np.load(velfile)
global nframes, npart, ndim, pos, vel
def init_minimizer(cadence):
pos = pos_full[::cadence]
vel = vel_full[::cadence]
def update_t(self, t):
bar_pot = agama.Potential(file=self.ai_bar(t))
dark_pot = agama.Potential(file=self.ai_dark(t))
self.potential = agama.Potential(bar_pot, dark_pot)
self.af = agama.ActionFinder(self.potential, interp=False)
def fitPotential(sim_name,
nsnap=600,
symmetry='a',
subsample_factor=1,
rmax_sel=600,
rmax_ctr=10,
rmax_exp=500,
save_coords=True):
'''
constructs a hybrid two-component basis expansion model of the potential for one Gizmo snapshot.
dark matter and hot gas are represented by an expansion in spherical harmonics.
remaining baryons (stars and cold gas) are represented by an azimuthal harmonic expansion in phi and a quintic spline in (R,z).
(see Agama docs, sections 2.2.2 and 2.2.3 for more details).
Arguments:
sim_name [str]: name of simulation folder within main dir
nsnap [int]: snapshot number
symmetry [char]: 's' for spherical, 'a' for axisymmetric, 't' for triaxial, 'n' for none (see table 2 of Agama docs)
subsample_factor [int]: factor by which to subsample (for speedup/testing)
rmax_sel [float]: max radius in kpc to select particles for fitting the model
rmax_ctr [float]: radius (kpc) defining the subset of stars which are used to define the reference frame (centering and rotation)
rmax_exp [float]: max radial extent in kpc of the hybrid potential expansion (both components)
save_coords [bool]: save center position, velocity, mean acceleration, and rotation matrix for principal-axis frame, to a file
'''
print('reading snapshot')
#read in the snapshot
part = ga.io.Read.read_snapshots(
species=['gas', 'star', 'dark'],
snapshot_values=nsnap,
simulation_directory=sims_dir + sim_name,
# snapshot_directory='output_accel',
particle_subsample_factor=subsample_factor,
assign_host_coordinates=True,
assign_host_principal_axes=True)
# start with default centering and rotation to define aperture
dist = ut.particle.get_distances_wrt_center(
part,
species=['gas', 'star', 'dark'],
rotation=part.host_rotation_tensors[0],
total_distance=True)
dist_vectors = ut.particle.get_distances_wrt_center(
part,
species=['gas', 'star', 'dark'],
rotation=part.host_rotation_tensors[0])
# compute new centering and rotation using a fixed aperture in stars
sp = 'star'
ctr_indices = np.where(dist[sp] < rmax_ctr)[0]
m = part[sp]['mass'][ctr_indices]
pos = part[sp]['position'][ctr_indices]
vel = part[sp]['velocity'][ctr_indices]
new_ctr = np.multiply(m, pos.T).sum(axis=1) / np.sum(m)
new_vctr = np.multiply(m, vel.T).sum(axis=1) / np.sum(m)
new_rot = ut.particle.get_principal_axes(part,
'star',
part_indices=ctr_indices,
print_results=False)
#optionally compute acceleration of center of mass frame if it was recorded
save_accel = ('acceleration' in part[sp].keys())
if save_accel:
print('saving acceleration of COM frame')
accel = part[sp]['acceleration'][ctr_indices]
new_actr = np.multiply(m, accel.T).sum(axis=1) / np.sum(m)
# recompute distances in the new frame
dist = ut.particle.get_distances_wrt_center(
part,
species=['star', 'gas', 'dark'],
center_position=new_ctr,
rotation=new_rot['rotation.tensor'],
total_distance=True)
dist_vectors = ut.particle.get_distances_wrt_center(
part,
species=['star', 'gas', 'dark'],
center_position=new_ctr,
rotation=new_rot['rotation.tensor'])
#pick out gas and stars within the region that we want to supply to the model
m_gas_tot = part['gas']['mass'].sum() * subsample_factor
pos_pa_gas = dist_vectors['gas'][dist['gas'] < rmax_sel]
m_gas = part['gas']['mass'][dist['gas'] < rmax_sel] * subsample_factor
print('{0:.3g} of {1:.3g} solar masses in gas selected'.format(
m_gas.sum(), m_gas_tot))
m_star_tot = part['star']['mass'].sum() * subsample_factor
pos_pa_star = dist_vectors['star'][dist['star'] < rmax_sel]
m_star = part['star']['mass'][dist['star'] < rmax_sel] * subsample_factor
print('{0:.3g} of {1:.3g} solar masses in stars selected'.format(
m_star.sum(), m_star_tot))
#separate cold gas in disk (modeled with cylspline) from hot gas in halo (modeled with multipole)
tsel = (np.log10(part['gas']['temperature']) < 4.5)
rsel = (dist['gas'] < rmax_sel)
pos_pa_gas_cold = dist_vectors['gas'][tsel & rsel]
m_gas_cold = part['gas']['mass'][tsel & rsel] * subsample_factor
print(
'{0:.3g} of {1:.3g} solar masses are cold gas to be modeled with cylspline'
.format(m_gas_cold.sum(), m_gas.sum()))
pos_pa_gas_hot = dist_vectors['gas'][(~tsel) & rsel]
m_gas_hot = part['gas']['mass'][(~tsel) & rsel] * subsample_factor
print(
'{0:.3g} of {1:.3g} solar masses are hot gas to be modeled with multipole'
.format(m_gas_hot.sum(), m_gas.sum()))
#combine components that will be fed to the cylspline part
pos_pa_bar = np.vstack((pos_pa_star, pos_pa_gas_cold))
m_bar = np.hstack((m_star, m_gas_cold))
#pick out the dark matter
m_dark_tot = part['dark']['mass'].sum() * subsample_factor
rsel = dist['dark'] < rmax_sel
pos_pa_dark = dist_vectors['dark'][rsel]
m_dark = part['dark']['mass'][rsel] * subsample_factor
print('{0:.3g} of {1:.3g} solar masses in dark matter selected'.format(
m_dark.sum(), m_dark_tot))
#stack with hot gas for multipole density
pos_pa_dark = np.vstack((pos_pa_dark, pos_pa_gas_hot))
m_dark = np.hstack((m_dark, m_gas_hot))
if save_coords:
#save the Hubble parameter to transform to comoving units
hubble = part.info['hubble']
scalefactor = part.info['scalefactor']
del (part)
#right now, configured to save to a new directory in the simulation directory.
#Recommended, since it's generally useful to have around
output_stem = sims_dir + sim_name + '/potential/{0:.0f}kpc/{1}_{2:d}'.format(
rmax_ctr, sim_name, nsnap)
try: # create the directory if it didn't exist
os.makedirs(os.path.dirname(output_stem))
except OSError as e:
if e.errno != errno.EEXIST:
raise
if save_coords:
cname = '{0}_coords.txt'.format(output_stem)
print('Saving coordinate transformation to {0}'.format(cname))
with open(cname, 'w') as f:
f.write(
'# Hubble parameter and scale factor (to convert physical <-> comoving) \n'
)
f.write('{0:.18g} {1:.18g}\n'.format(hubble, scalefactor))
f.write('# center position (kpc comoving)\n')
np.savetxt(f, new_ctr)
f.write('# center velocity (km/s physical)\n')
np.savetxt(f, new_vctr)
if save_accel:
f.write('# center acceleration (km/s^2 physical)\n')
np.savetxt(f, new_actr)
f.write('# rotation to principal-axis frame\n')
np.savetxt(f, new_rot['rotation.tensor'])
print(
'Computing multipole expansion coefficients for dark matter/hot gas component'
)
p_dark = agama.Potential(type='multipole',
particles=(pos_pa_dark, m_dark),
lmax=4,
symmetry=symmetry,
rmin=0.1,
rmax=rmax_exp)
p_dark.export('{0}.dark.{1}.coef_mul'.format(output_stem,
symmlabel[symmetry]))
print(
'Computing cylindrical spline coefficients for stellar/cold gas component'
)
p_bar = agama.Potential(
type='cylspline',
particles=(pos_pa_bar, m_bar),
mmax=4,
symmetry=symmetry,
#gridsizer=40, gridsizez=40,
rmin=0.1,
rmax=rmax_exp)
p_bar.export('{0}.bar.{1}.coef_cylsp'.format(output_stem,
symmlabel[symmetry]))
print('done, enjoy your potentials!')
def update_index(self, index, ss_id=None, snap=None):
agama.setUnits(mass=1, length=1, velocity=1)
self.current_index = index
self.ss_id = ss_id
if snap is not None:
self.snap = snap
else:
self.snap = \
gizmo.io.Read.read_snapshots(['star', 'gas', 'dark'],
'index', index,
properties=['id', 'position',
'velocity', 'mass',
'form.scalefactor'],
simulation_directory=
self.simulation_directory,
assign_principal_axes=True)
if self.sim_name is None:
head = gizmo.io.Read.read_header(snapshot_value=self.startnum,
simulation_directory=
self.simulation_directory)
self.sim_name = head['simulation.name'].replace(" ", "_")
potential_cache_file = self.cache_directory + '/potential_id'+str(index)
potential_cache_file += '_' + self.sim_name + '_pot'
try:
self.potential = agama.Potential(file=potential_cache_file)
except:
star_position = self.snap['star'].prop('host.distance.principal')
gas_position = self.snap['gas'].prop('host.distance.principal')
dark_position = self.snap['dark'].prop('host.distance.principal')
star_mass = self.snap['star']['mass']
gas_mass = self.snap['gas']['mass']
dark_mass = self.snap['dark']['mass']
position = np.concatenate((star_position, gas_position))
mass = np.concatenate((star_mass, gas_mass))
#TODO make these user-controllable
self.pdark = agama.Potential(type="Multipole",
particles=(dark_position, dark_mass),
symmetry='a', gridsizeR=20, lmax=2)
self.pbar = agama.Potential(type="CylSpline",
particles=(position, mass),
gridsizer=20, gridsizez=20,
mmax=0, Rmin=0.2, symmetry='a',
Rmax=50, Zmin=0.02, Zmax=10)
self.potential = agama.Potential(self.pdark, self.pbar)
self.potential.export(potential_cache_file)
if ss_id is not None:
self.ss_init = True
ss_key = np.where(self.snap['star']['id'] == self.ss_id)[0]
self.chosen_position = self.snap['star'].prop(
'host.distance.principal')[ss_key]
self.chosen_velocity = self.snap['star'].prop(
'host.velocity.principal')[ss_key]
else:
self.ss_init = False
self.af = agama.ActionFinder(self.potential, interp=False)
def __init__(self, filename):
'''
Initialize starting values of scaled parameters and and define upper/lower limits on them;
also obtain the true values by analyzing the input file name
'''
self.initValues = numpy.array(
[9.0, 0.1, 0.5, 4.0, 1.0, 6.0, 1.5, 1.0, 1.0, 1.0, 1.])
self.minValues = numpy.array(
[6.0, -1.0, -0.5, 2.2, 0.5, 3.2, 0.0, 0.5, 0.1, 0.1, -1.])
self.maxValues = numpy.array(
[10., 2.0, 1.9, 6.0, 2.5, 12., 2.8, 2.0, 2.8, 2.8, 3.])
self.labels = ( \
r'$\log_{10}(\rho_\mathrm{scale})$', r'$\log_{10}(R_\mathrm{scale})$', r'$\gamma$', r'$\beta$', r'$\alpha$', \
r'$B_\mathrm{DF}$', r'$\Gamma_\mathrm{DF}$', r'$\eta$', r'$g_r$', r'$h_r$', r'$\log_{10}(J_0)$')
self.numPotParams = 5 # potential params come first in the list
self.numDFParams = 6 # DF params come last in the list
# parse true values
m = re.search(
r'gs(\d+)_bs(\d+)_rcrs([a-z\d]+)_rarc([a-z\d]+)_([a-z]+)_(\d+)mpc3',
filename)
n = re.search(r'data_([ch])_rh(\d+)_rs(\d+)_gs(\d+)', filename)
if m:
self.trueParams = (
numpy.log10(float(m.group(6)) * 1e6),
0.0,
1.0 if m.group(5) == 'cusp' else 0.0,
3.0,
1.0,
# true params of tracer DF do not belong to this family of models, so ignore them
numpy.nan,
numpy.nan,
numpy.nan,
numpy.nan,
numpy.nan,
numpy.nan)
tracerParams = dict(type='Spheroid',
densityNorm=1.0,
scaleRadius=float(m.group(3)) * 0.01,
gamma=float(m.group(1)) * 0.01,
beta=float(m.group(2)) * 0.1,
alpha=2.0)
beta0 = 0
r_a = float(m.group(4)) * 0.01 * float(
m.group(3)) * 0.01 if m.group(4) != 'inf' else numpy.inf
elif n:
self.trueParams = (
numpy.log10(3.021516e7 if n.group(1) == 'c' else 2.387329e7),
numpy.log10(float(n.group(2))),
0.0 if n.group(1) == 'c' else 1.0, 4.0, 1.0, numpy.nan,
numpy.nan, numpy.nan, numpy.nan, numpy.nan, numpy.nan)
tracerParams = dict(
type='Spheroid',
densityNorm=1.0,
scaleRadius=1.75 if n.group(3) == '175' else 0.5,
gamma=float(n.group(4)) * 0.1,
beta=5.0,
alpha=2.0)
beta0 = -0.5
r_a = numpy.inf
else:
print("Can't determine true parameters!")
return
self.truePotential = agama.Potential(
type='Spheroid',
densityNorm=10**self.trueParams[0],
scaleRadius=10**self.trueParams[1],
gamma=self.trueParams[2],
beta=self.trueParams[3],
alpha=self.trueParams[4])
# normalize the tracer density profile to have a total mass of unity
tracerParams["densityNorm"] /= agama.Density(
**tracerParams).totalMass()
self.tracerDensity = agama.Density(**tracerParams)
self.tracerBeta = lambda r: (beta0 + (r / r_a)**2) / (1 + (r / r_a)**2)
