def rewind_zvz(z_init, vz_init, mass, omega, step):
grewind = wendypy.nbody(z_init,
vz_init,
mass,
-step,
omega=omega,
approx=True,
nleap=1)
z_rewind, vz_rewind = next(grewind)
return (z_rewind, vz_rewind)
def forward_nstep_zvz(z_init, vz_init, mass, omega, step, nstep):
gforward = wendypy.nbody(z_init,
vz_init,
mass,
-step,
omega=omega,
approx=True,
nleap=1)
for ii in range(nstep):
z_forward, vz_forward = next(gforward)
return (z_forward, vz_forward)
def rewind_nstep_zvz(z_init, vz_init, mass, omega, step, nstep):
grewind = wendypy.nbody(z_init,
vz_init,
mass,
-step,
omega=omega,
approx=True,
nleap=1)
for ii in range(nstep):
z_rewind, vz_rewind = next(grewind)
return (z_rewind, vz_rewind)
def zvzdiff(z_init, vz_init, mass, omega1, omega2, step):
# (temporary?) way to deal with small masses
relevant_particles_index = mass > (numpy.median(mass[mass > 10.**-9.]) *
10.**-6.)
if numpy.any(mass[relevant_particles_index] < (
10.**-8. * numpy.median(mass[relevant_particles_index]))):
print(
numpy.sum(mass[relevant_particles_index] < (
10.**-8. * numpy.median(mass[relevant_particles_index]))))
# integrate with wendy
g1 = wendypy.nbody(z_init[relevant_particles_index],
vz_init[relevant_particles_index],
mass[relevant_particles_index],
step,
omega=omega1,
approx=True,
nleap=1)
z_next1, vz_next1 = next(g1)
dz1 = numpy.zeros_like(z_init)
dvz1 = numpy.zeros_like(z_init)
dz1[relevant_particles_index] = z_next1 - z_init[relevant_particles_index]
dvz1[relevant_particles_index] = vz_next1 - vz_init[
relevant_particles_index]
g2 = wendypy.nbody(z_init[relevant_particles_index],
vz_init[relevant_particles_index],
mass[relevant_particles_index],
step,
omega=omega2,
approx=True,
nleap=1)
z_next2, vz_next2 = next(g2)
dz2 = numpy.zeros_like(z_init)
dvz2 = numpy.zeros_like(z_init)
dz2[relevant_particles_index] = z_next2 - z_init[relevant_particles_index]
dvz2[relevant_particles_index] = vz_next2 - vz_init[
relevant_particles_index]
return (dz2 - dz1, dvz2 - dvz1)
def fit_m2m(w_init,
z_init,
vz_init,
omega_m2m,
zsun_m2m,
data_dicts,
step=0.001,
nstep=1000,
eps=0.1,
mu=1.,
prior='entropy',
w_prior=None,
kernel=hom2m.epanechnikov_kernel,
kernel_deriv=hom2m.epanechnikov_kernel_deriv,
h_m2m=0.02,
npop=1,
smooth=None,
st96smooth=False,
output_wevolution=False,
output_zvzevolution=False,
fit_zsun=False,
fit_omega=False,
skipomega=10,
delta_omega_frac=0.1,
number_density=False,
xnm_m2m=1.0,
skipxnm=10,
fit_xnm=False):
"""
NAME:
fit_m2m
PURPOSE:
Run M2M optimization for wendy M2M
INPUT:
w_init - initial weights [N] or [N,npop]
z_init - initial z [N]
vz_init - initial vz (rad) [N]
omega_m2m - background potential parameter omega, if None no background
zsun_m2m - Sun's height above the plane [N]
data_dicts - list of dictionaries that hold the data, these are described in more detail below
step= stepsize of orbit integration
nstep= number of steps to integrate the orbits for
eps= M2M epsilon parameter (can be array when fitting zsun, omega; in that case eps[0] = eps_weights, eps[1] = eps_zsun, eps[1 or 2 based on fit_zsun] = eps_omega)
mu= M2M entropy parameter mu
prior= ('entropy' or 'gamma')
w_prior= (None) prior weights (if None, equal to w_init)
fit_zsun= (False) if True, also optimize zsun
fit_omega= (False) if True, also optimize omega
skipomega= only update omega every skipomega steps
delta_omega_frac= (0.1) fractional difference in omega to use to compute derivative of objective function wrt omega
kernel= a smoothing kernel
kernel_deriv= the derivative of the smoothing kernel
h_m2m= kernel size parameter for computing the observables
npop= (1) number of theoretical populations
smooth= smoothing parameter alpha (None for no smoothing)
st96smooth= (False) if True, smooth the constraints (Syer & Tremaine 1996), if False, smooth the objective function and its derivative (Dehnen 2000)
output_wevolution= if set to an integer, return the time evolution of this many randomly selected weights
output_zvzevolution= if set to an integer, return the time evolution
of this many randomly selected weights only when
output_wevolution is True
number_density = (False) if True, observed density is number density and density calculation requires xnm
xnm_m2m = (1.0) initial value of xnm: number_density/mass_density [1], assuming a single population
skipxnm = only update Xnm every skipxnm steps
fit_xnm = (False) if True, also optimise xnm
DATA DICTIONARIES:
The data dictionaries have the following form:
'type': type of measurement: 'dens', 'v2'
'pops': the theoretical populations included in this measurement;
single number or list
'zobs': vertical height of the observation
'zrange': width of vertical bin relative to some fiducial value (used to scale h_m2m, which should therefore be appropriate for the fiducial value)
'obs': the actual observation
'unc': the uncertainty in the observation
of these, zobs, obs, and unc can be arrays for mulitple measurements
OUTPUT:
(w_out,[zsun_out, [omega_out, [xnm_out]],z_m2m,vz_m2m,Q_out,[wevol,rndindx]) -
(output weights [N],
[Solar offset [nstep] optional],
[omega [nstep] optional when fit_omega],
z_m2m [N] final z,
vz_m2m [N] final vz,
objective function as a function of time [nstep],
[weight evolution for randomly selected weights,index of random weights])
HISTORY:
2017-07-20 - Started from hom2m.fit_m2m - Bovy (UofT)
2018-10-16 - add external potential using omega_m2m - Kawata (MSSL/UCL)
2018-10-29 - add xnm=number_ensity/mass_density - Kawata (MSSL/UCL)
"""
if len(w_init.shape) == 1:
w_out = numpy.empty((len(w_init), npop))
w_out[:, :] = numpy.tile(copy.deepcopy(w_init), (npop, 1)).T
else:
w_out = copy.deepcopy(w_init)
zsun_out = numpy.empty(nstep)
omega_out = numpy.empty(nstep)
if number_density:
xnm_out = numpy.empty(nstep)
else:
xnm_out = numpy.ones(nstep)
xnm_m2m = 1.0
if w_prior is None:
if len(w_init.shape) == 1:
w_prior = numpy.empty((len(w_init), npop))
w_prior[:, :] = numpy.tile(copy.deepcopy(w_init), (npop, 1)).T
else:
w_piror = copy.deepcopy(w_init)
else:
if len(w_prior.shape) == 1:
w_prior = numpy.tile(w_prior, (npop, 1)).T
# Parse data_dict
data_dict = parse_data_dict(data_dicts)
# Parse eps
if isinstance(eps, float):
eps = [eps]
if fit_zsun: eps.append(eps[0])
if fit_omega: eps.append(eps[0])
if fit_xnm: eps.append(eps[0])
Q_out = []
if output_wevolution:
rndindx = numpy.random.permutation(len(w_out))[:output_wevolution]
wevol = numpy.zeros((output_wevolution, npop, nstep))
if output_zvzevolution:
zevol = numpy.zeros((output_wevolution, nstep))
vzevol = numpy.zeros((output_wevolution, nstep))
# Compute force of change for first iteration
fcw, delta_m2m_new = \
force_of_change_weights(w_out,zsun_m2m,z_init,vz_init,
data_dicts,prior,mu,w_prior,
h_m2m=h_m2m,kernel=kernel,
xnm_m2m=xnm_m2m)
fcw *= w_out
fcz = 0.
if fit_zsun:
fcz = force_of_change_zsun(w_init,
zsun_m2m,
z_init,
vz_init,
z_obs,
dens_obs_noise,
delta_m2m_new,
densv2_obs_noise,
deltav2_m2m_new,
kernel=kernel,
kernel_deriv=kernel_deriv,
h_m2m=h_m2m)
fcxnm = 0.0
if fit_xnm:
fcxnm = force_of_change_xnm(w_out,
zsun_m2m,
z_init,
vz_init,
data_dicts,
delta_m2m_new,
h_m2m=h_m2m,
kernel=kernel,
xnm_m2m=xnm_m2m)
fco = 0.0
# Rewind for first step
mass = numpy.sum(w_out, axis=1)
z_prev, vz_prev = rewind_zvz(z_init, vz_init, mass, omega_m2m, step)
if fit_omega:
fco = force_of_change_omega(w_out,
zsun_m2m,
omega_m2m,
z_init,
vz_init,
z_prev,
vz_prev,
step,
data_dicts,
delta_m2m_new,
h_m2m=h_m2m,
kernel=kernel,
delta_omega_frac=delta_omega_frac,
xnm_m2m=xnm_m2m)
if not smooth is None:
delta_m2m = delta_m2m_new
else:
delta_m2m = [None for d in data_dicts]
if not smooth is None and not st96smooth:
Q = [d**2 for d in delta_m2m**2.]
# setup skipomega omega counter and prev. (z,vz) for F(omega)
# xcounter= skipxnm-1 # Causes F(Xnm) to be computed in the 1st step
# ocounter= skipomega-1 # Causes F(omega) to be computed in the 1st step
# no update for omega or xnm at the first step.
xcounter = 0
ocounter = 0
z_m2m, vz_m2m = z_init, vz_init
for ii in range(nstep):
# Update weights first
if True:
w_out += eps[0] * step * fcw
w_out[w_out < 10.**-16.] = 10.**-16.
# then zsun
if fit_zsun:
zsun_m2m += eps[1] * step * fcz
zsun_out[ii] = zsun_m2m
# then xnm
if fit_xnm and xcounter == skipxnm:
# assume single population and use fcxnm[0] only
dxnm = eps[1 + fit_zsun + fit_omega] * step * fcxnm[0]
# print(' xnm_m2m, dxnm=',xnm_m2m,dxnm)
max_dxnm = xnm_m2m / 10.0
if numpy.fabs(dxnm) > max_dxnm:
dxnm = max_dxnm * numpy.sign(dxnm)
xnm_m2m += dxnm
# print(ii,' step Xnm=',xnm_m2m, eps[1+fit_zsun+fit_omega])
xcounter = 0
# then omega (skipped in the first step, so undeclared vars okay)
if fit_omega and ocounter == skipomega:
domega = eps[1 + fit_zsun] * step * fco
# print(' omega_m2m, dom=',omega_m2m,domega, numpy.sum(delta_m2m_new, axis=1))
# max_domega= delta_omega/30.
max_domega = omega_m2m / 10.0
if numpy.fabs(domega) > max_domega:
domega = max_domega * numpy.sign(domega)
omega_m2m += domega
# Keep (z,vz) the same in new potential
# A_now, phi_now= zvz_to_Aphi(z_m2m,vz_m2m,omega_m2m)
ocounter = 0
# (Store objective function)
if not smooth is None and st96smooth:
Q_out.append([d**2. for d in delta_m2m])
elif not smooth is None:
Q_out.append(copy.deepcopy(Q))
else:
Q_out.append(
[d**2. for d in list(chain.from_iterable(delta_m2m_new))])
# Then update the dynamics
mass = numpy.sum(w_out, axis=1)
# (temporary?) way to deal with small masses
relevant_particles_index = mass > (
numpy.median(mass[mass > 10.**-9.]) * 10.**-6.)
if numpy.any(mass[relevant_particles_index] < (
10.**-8. * numpy.median(mass[relevant_particles_index]))):
print(
numpy.sum(mass[relevant_particles_index] < (
10.**-8. * numpy.median(mass[relevant_particles_index]))))
# g= wendypy.nbody(z_m2m[relevant_particles_index],
# vz_m2m[relevant_particles_index],
# mass[relevant_particles_index],
# step, omega=omega_m2m, maxcoll=10000000)
g = wendypy.nbody(z_m2m[relevant_particles_index],
vz_m2m[relevant_particles_index],
mass[relevant_particles_index],
step,
omega=omega_m2m,
approx=True,
nleap=1)
tz_m2m, tvz_m2m = next(g)
z_m2m[relevant_particles_index] = tz_m2m
vz_m2m[relevant_particles_index] = tvz_m2m
z_m2m -= numpy.sum(mass * z_m2m) / numpy.sum(mass)
vz_m2m -= numpy.sum(mass * vz_m2m) / numpy.sum(mass)
# Compute force of change
if smooth is None or not st96smooth:
# Turn these off
tdelta_m2m = None
else:
tdelta_m2m = delta_m2m
fcw_new, delta_m2m_new = \
force_of_change_weights(w_out,zsun_m2m,z_m2m,vz_m2m,
data_dicts,prior,mu,w_prior,
h_m2m=h_m2m,kernel=kernel,
delta_m2m=tdelta_m2m, xnm_m2m=xnm_m2m)
fcw_new *= w_out
if fit_zsun:
if smooth is None or not st96smooth:
tdelta_m2m = delta_m2m_new
fcz_new = force_of_change_zsun(w_out,
zsun_m2m,
z_m2m,
vz_m2m,
data_dicts,
tdelta_m2m,
kernel=kernel,
kernel_deriv=kernel_deriv,
h_m2m=h_m2m)
if fit_xnm:
# Update Xnm in this step?
xnm_out[ii] = xnm_m2m
xcounter += 1
if xcounter == skipxnm:
if smooth is None or not st96smooth:
tdelta_m2m = delta_m2m_new
fcxnm_new = force_of_change_xnm(w_out,
zsun_m2m,
z_m2m,
vz_m2m,
data_dicts,
tdelta_m2m,
h_m2m=h_m2m,
kernel=kernel,
xnm_m2m=xnm_m2m)
if fit_omega:
omega_out[ii] = omega_m2m
# Update omega in this step?
ocounter += 1
if ocounter == skipomega:
if not fit_zsun and (smooth is None or not st96smooth):
tdelta_m2m = delta_m2m_new
# tdeltav2_m2m= deltav2_m2m_new
# use step to evaluate gradient of omega
fco_new = force_of_change_omega(
w_out,
zsun_m2m,
omega_m2m,
z_m2m,
vz_m2m,
z_prev,
vz_prev,
step,
data_dicts,
tdelta_m2m,
h_m2m=h_m2m,
kernel=kernel,
delta_omega_frac=delta_omega_frac,
xnm_m2m=xnm_m2m)
# store previous position and velocity every step.
z_prev = copy.copy(z_m2m)
vz_prev = copy.copy(vz_m2m)
# Increment smoothing
if not smooth is None and st96smooth:
delta_m2m = [
d + step * smooth * (dn - d)
for d, dn in zip(list(chain.from_iterable(delta_m2m)),
list(chain.from_iterable(delta_m2m_new)))
]
fcw = fcw_new
if fit_zsun: fcz = fcz_new
if fit_omega and ocounter == skipomega: fco = fco_new
if fit_xnm and xcounter == skipxnm: fcxnm = fcxnm_new
elif not smooth is None:
Q_new = [d**2. for d in delta_m2m_new]
Q = [q + step * smooth * (qn - q) for q, qn in zip(Q, Q_new)]
fcw += step * smooth * (fcw_new - fcw)
if fit_zsun: fcz += step * smooth * (fcz_new - fcz)
if fit_xnm and xcounter == skipxnm:
fcxnm += step * smooth * (fcxnm - fcxnm_new)
if fit_omega and ocounter == skipomega:
fco += step * smooth * (fco_new - fco)
else:
fcw = fcw_new
if fit_zsun: fcz = fcz_new
if fit_xnm and xcounter == skipxnm: fcxnm = fcxnm_new
if fit_omega and ocounter == skipomega: fco = fco_new
# Record random weights if requested
if output_wevolution:
wevol[:, :, ii] = w_out[rndindx]
if output_zvzevolution:
zevol[:, ii] = z_m2m[rndindx]
vzevol[:, ii] = vz_m2m[rndindx]
out = (w_out, )
if fit_zsun: out = out + (zsun_out, )
if fit_omega:
out = out + (omega_out, )
if fit_xnm:
out = out + (xnm_out, )
out = out + (
z_m2m,
vz_m2m,
)
out = out + (numpy.array(Q_out), )
if output_wevolution:
out = out + (
wevol,
rndindx,
)
if output_zvzevolution:
out = out + (zevol, )
out = out + (vzevol, )
return out
def force_of_change_omega(w_m2m,
zsun_m2m,
omega_m2m,
z_m2m,
vz_m2m,
z_prev,
vz_prev,
step,
data_dicts,
delta_m2m,
h_m2m=0.02,
kernel=hom2m.epanechnikov_kernel,
delta_omega_frac=0.1,
xnm_m2m=1.0):
"""Compute the force of change by direct finite difference
of the objective function"""
mass = numpy.sum(w_m2m, axis=1)
# delta omega to evaluate the gradient
delta_omega = omega_m2m * delta_omega_frac
# difference in dz and dvz due to the different omega
# dz, dvz = zvzdiff(
# z_prev, vz_prev, mass, omega_m2m, omega_m2m+delta_omega, step)
# fcw, delta_m2m_do = force_of_change_weights(\
# w_m2m,zsun_m2m,z_m2m+dz,vz_m2m+dvz,
# data_dicts,
# 'entropy',0.,1., # weights prior doesn't matter, so set to zero
# h_m2m=h_m2m,kernel=kernel,
# delta_m2m=delta_m2m, xnm_m2m=xnm_m2m)
# return -numpy.nansum(\
# delta_m2m[0]*(delta_m2m_do[0]-delta_m2m[0])
# +delta_m2m[1]*(delta_m2m_do[1]-delta_m2m[1]))\
# /delta_omega
# more directly computing derivative, with no mass limit
omega1 = omega_m2m
omega2 = omega_m2m + delta_omega
# integrate with wendy with omega1
g1 = wendypy.nbody(z_prev, vz_prev, mass, step, omega=omega1, \
approx=True, nleap=1)
z_next1, vz_next1 = next(g1)
# evaluate delta
fcw, delta_m2m_do1 = force_of_change_weights(\
w_m2m,zsun_m2m,z_next1,vz_next1,
data_dicts,
'entropy',0.,1., # weights prior doesn't matter, so set to zero
h_m2m=h_m2m,kernel=kernel,
delta_m2m=delta_m2m, xnm_m2m=xnm_m2m)
# integrate with omega2
g2 = wendypy.nbody(z_prev, vz_prev, mass, step, omega=omega2, \
approx=True, nleap=1)
z_next2, vz_next2 = next(g2)
fcw, delta_m2m_do2 = force_of_change_weights(\
w_m2m,zsun_m2m,z_next2,vz_next2,
data_dicts,
'entropy',0.,1., # weights prior doesn't matter, so set to zero
h_m2m=h_m2m,kernel=kernel,
delta_m2m=delta_m2m, xnm_m2m=xnm_m2m)
return -numpy.nansum(\
delta_m2m_do1[0]*(delta_m2m_do2[0]-delta_m2m_do1[0])
+delta_m2m_do1[1]*(delta_m2m_do2[1]-delta_m2m_do2[1]))\
/delta_omega
h_def = 0.1 # default h
print(' disk parameters, np, sigma, Mtot, omega, zsun, Xnm=', \
n_init, sigma_true, totmass_true, omegadm_true, zsun_true, xnm_true)
print(' smoothing length (default) =', h_def)
n_init0 = n_init
zh_true = sigma_true**2. / totmass_true # Where 2\pi G= 1 so units of zh are ~311 pc
tdyn = zh_true / sigma_true
z_init, vz_init, m_init = wendym2m.sample_sech2(sigma_true,
totmass_true,
n=n_init)
print('zh, tdyn, omega_dm =', zh_true, tdyn, omegadm_true)
# run Wendy to introduce the background potential adiabatically
print(' running Wendy to add the background potential adiabatically')
g = wendypy.nbody(z_init, vz_init, m_init, dttdyn * tdyn, approx=True, nleap=1)
nt = 4000
zt = numpy.empty((n_init, nt + 1))
vzt = numpy.empty((n_init, nt + 1))
# Et = numpy.empty((nt+1))
zt[:, 0] = z_init
vzt[:, 0] = vz_init
# Et[0] = wendy.energy(z_init, vz_init, m_init)
# increasing omega
nstep_omega = 1000
if nstep_omega > nt:
print(' Error nstep_omega=', nstep_omega, ' should be larger than nt=', nt)
sys.exit()
domega = omegadm_true / nstep_omega
omega_ii = 0.0
tz = z_init
