def test_againstexact_sech2disk_manyparticles(): # Test that the exact N-body and the approximate N-body agree N = 101 totmass = 1. sigma = 1. zh = 2. * sigma**2. / totmass x = numpy.arctanh(2. * numpy.random.uniform(size=N) - 1) * zh v = numpy.random.normal(size=N) * sigma v -= numpy.mean(v) # stabilize m = numpy.ones_like(x) / N * (1. + 0.1 * (2. * numpy.random.uniform(size=N) - 1)) g = wendy.nbody(x, v, m, 0.05, approx=True, nleap=2000) ge = wendy.nbody(x, v, m, 0.05) cnt = 0 while cnt < 100: tx, tv = next(g) txe, tve = next(ge) assert numpy.all( numpy.fabs(tx - txe) < 10.**-5. ), "Exact and approximate N-body give different positions" assert numpy.all( numpy.fabs(tv - tve) < 10.**-5. ), "Exact and approximate N-body give different positions" cnt += 1 return None
def test_notracermasses_error(): # Test that an exception is raised when some of the masses are zero # (or just tiny) x= numpy.array([-1.,1.]) v= numpy.array([0.,0.]) # Test zero m= numpy.array([1.,0.]) g= wendy.nbody(x,v,m,2) with pytest.raises(ValueError) as excinfo: tx,tv= next(g) assert str(excinfo.value) == 'Tiny masses m much smaller than the median mass are not supported, please remove these from the inputs' # Also test tiny m= numpy.array([1.,10.**-10.]) g= wendy.nbody(x,v,m,2) with pytest.raises(ValueError) as excinfo: tx,tv= next(g) assert str(excinfo.value) == 'Tiny masses m much smaller than the median mass are not supported, please remove these from the inputs' return None
def test_maxncoll_error(): # Simple test where we know the number of collisions x= numpy.array([-1.,1.]) v= numpy.array([0.,0.]) m= numpy.array([1.,1.]) # First collision at t=sqrt(2) g= wendy.nbody(x,v,m,2,maxcoll=0,full_output=True) with pytest.raises(RuntimeError) as excinfo: tx,tv,ncoll, _= next(g) assert str(excinfo.value) == 'Maximum number of collisions per time step exceeded' return None
def test_omegadt_lt_pi2(): # Test that wendy.nbody raises a ValueError when omega x dt > pi/2 omega = 1. + 10.**-8. dt = numpy.pi / 2. with pytest.raises(ValueError) as excinfo: g = wendy.nbody(None, None, None, dt, omega=omega) next(g) assert str( excinfo.value ) == 'When omega is set, omega*dt needs to be less than pi/2; please adjust dt' return None
def test_nleap_error(): # Code should raise ValueError if nleap is not specified for approx. calc. x = numpy.array([-1., 1.]) v = numpy.array([0., 0.]) m = numpy.array([1., 1.]) g = wendy.nbody(x, v, m, 2, approx=True) with pytest.raises(ValueError) as excinfo: tx, tv, ncoll, _ = next(g) assert str( excinfo.value ) == 'When approx is True, the number of leapfrog steps nleap= per output time step needs to be set' return None
def test_maxncoll_warn(): # Simple test where we know the number of collisions x= numpy.array([-1.,1.]) v= numpy.array([0.,0.]) m= numpy.array([1.,1.]) # First collision at t=sqrt(2) g= wendy.nbody(x,v,m,2,maxcoll=0,full_output=True,warn_maxcoll=True) with pytest.warns(RuntimeWarning) as record: tx,tv,ncoll, _= next(g) # check that only one warning was raised assert len(record) == 1 # check that the message matches assert record[0].message.args[0] == "Maximum number of collisions per time step exceeded" return None
def test_momentum_conservation_unequalmasses(): # Test that momentum is conserved for a simple problem x= numpy.array([-1.1,0.1,1.3]) v= numpy.array([3.,2.,-5.]) m= numpy.array([1.,2.,3.]) g= wendy.nbody(x,v,m,0.05) p= wendy.momentum(v,m) cnt= 0 while cnt < 100: tx,tv= next(g) assert numpy.fabs(wendy.momentum(tv,m)-p) < 10.**-10., "Momentum not conserved during simple N-body integration" cnt+= 1 return None
def test_energy_conservation(): # Test that energy is conserved for a simple problem x= numpy.array([-1.1,0.1,1.3]) v= numpy.array([3.,2.,-5.]) m= numpy.array([1.,1.,1.]) g= wendy.nbody(x,v,m,0.05) E= wendy.energy(x,v,m) cnt= 0 while cnt < 100: tx,tv= next(g) assert numpy.fabs(wendy.energy(tx,tv,m)-E) < 10.**-10., "Energy not conserved during simple N-body integration" cnt+= 1 return None
def test_time(): # Just run the timer... N = 101 totmass = 1. sigma = 1. zh = 2. * sigma**2. / totmass x = numpy.arctanh(2. * numpy.random.uniform(size=N) - 1) * zh v = numpy.random.normal(size=N) * sigma v -= numpy.mean(v) # stabilize m = numpy.ones_like(x) / N * (1. + 0.1 * (2. * numpy.random.uniform(size=N) - 1)) g = wendy.nbody(x, v, m, 0.05, approx=True, nleap=1000, full_output=True) tx, tv, time_elapsed = next(g) assert time_elapsed < 1., 'More than 1 second elapsed for simple problem' return None
def test_energy_conservation_unequalmasses(): # Test that energy is conserved for a simple problem x = numpy.array([-1.1, 0.1, 1.3]) v = numpy.array([3., 2., -5.]) m = numpy.array([1., 2., 3.]) g = wendy.nbody(x, v, m, 0.05, approx=True, nleap=100000) E = wendy.energy(x, v, m) cnt = 0 while cnt < 100: tx, tv = next(g) assert numpy.fabs( wendy.energy(tx, tv, m) - E ) / E < 10.**-6., "Energy not conserved during approximate N-body integration" cnt += 1 return None
def test_total_energy(): # Test that total energy is sum of potential and kinetic energy x = numpy.array([-1.1, 0.1, 1.3]) v = numpy.array([3., 2., -5.]) m = numpy.array([1., 1., 1.]) g = wendy.nbody(x, v, m, 0.05) cnt = 0 while cnt < 10: tx, tv = next(g) Etot = wendy.energy(tx, tv, m, energy_type=wendy.EnergyType.TOTAL) Epot = wendy.energy(tx, tv, m, energy_type=wendy.EnergyType.POTENTIAL) Ekin = wendy.energy(tx, tv, m, energy_type=wendy.EnergyType.KINETIC) assert Etot == Epot + Ekin, "Total energy should be sum of kinetic and potential energies" cnt += 1 return None
def test_energy_ratio(): # Test that energy ratio gives ratio of kinetic to potential energy x = numpy.array([-1.1, 0.1, 1.3]) v = numpy.array([3., 2., -5.]) m = numpy.array([1., 1., 1.]) g = wendy.nbody(x, v, m, 0.05) cnt = 0 while cnt < 10: tx, tv = next(g) Erat = wendy.energy(tx, tv, m, energy_type=wendy.EnergyType.RATIO) Epot = wendy.energy(tx, tv, m, energy_type=wendy.EnergyType.POTENTIAL) Ekin = wendy.energy(tx, tv, m, energy_type=wendy.EnergyType.KINETIC) assert Erat == Ekin / Epot, "Energy ratio should be kinetic divided potential energy" cnt += 1 return None
def test_energy_conservation(): # Test that energy is conserved for a simple problem x = numpy.array([-1.1, 0.1, 1.3]) v = numpy.array([3., 2., -5.]) m = numpy.array([1., 1., 1.]) omega = 1.1 g = wendy.nbody(x, v, m, 0.05, omega=omega, approx=True, nleap=100000) E = wendy.energy(x, v, m, omega=omega) cnt = 0 while cnt < 100: tx, tv = next(g) assert numpy.fabs( wendy.energy(tx, tv, m, omega=omega) - E ) / E < 10.**-6., "Energy not conserved during approximate N-body integration with external harmonic potential" cnt += 1 return None
def test_density_profile(): # Test that density profile runs and gives reasonable results x = numpy.linspace(-5, 5, 100) v = x * 0.01 m = numpy.ones(x.shape) g = wendy.nbody(x, v, m, 0.05) cnt = 0 while cnt < 10: tx, tv = next(g) bin_centers, densities = wendy.density_profile(tx, m) assert len(bin_centers) == len(densities) assert (numpy.fabs(bin_centers) < 7).all(), "bin centers should not exceed range of x" assert (densities > 0).all() and (densities < 100).all( ), "densities should be positive and not too large" cnt += 1 return None
def test_energy_conservation_sech2disk_manyparticles(): # Test that energy is conserved for a self-gravitating disk N= 101 totmass= 1. sigma= 1. zh= 2.*sigma**2./totmass x= numpy.arctanh(2.*numpy.random.uniform(size=N)-1)*zh v= numpy.random.normal(size=N)*sigma v-= numpy.mean(v) # stabilize m= numpy.ones_like(x)/N*(1.+0.1*(2.*numpy.random.uniform(size=N)-1)) g= wendy.nbody(x,v,m,0.05) E= wendy.energy(x,v,m) cnt= 0 while cnt < 100: tx,tv= next(g) assert numpy.fabs(wendy.energy(tx,tv,m)-E) < 10.**-10., "Energy not conserved during simple N-body integration" cnt+= 1 return None
def test_samex(): # Test that the code works if two particles are at the exact same position # middle ones setup such that they end up in the same place for the first # force evaluation x = numpy.array( [-1.1, -2. * 0.05 / 2. / 10000, 0.3 * 0.05 / 2. / 10000, 1.3]) v = numpy.array([3., 2., -.3, -5.]) m = numpy.array([1., 1., 1., 1.]) g = wendy.nbody(x, v, m, 0.05, approx=True, nleap=10000) E = wendy.energy(x, v, m) cnt = 0 while cnt < 1: tx, tv = next(g) assert numpy.fabs( wendy.energy(tx, tv, m) - E ) / E < 10.**-6., "Energy not conserved during approximate N-body integration" cnt += 1 return None
def test_count_ncoll(): # Simple test where we know the number of collisions x = numpy.array([-1., 1.]) v = numpy.array([0., 0.]) m = numpy.array([1., 1.]) # First collision at t=sqrt(2) g = wendy.nbody(x, v, m, 1, full_output=True) tx, tv, ncoll, _ = next(g) assert ncoll == 0, 'Number of collisions in simple 2-body problem is wrong' tx, tv, ncoll, _ = next(g) # collision should have happened now assert ncoll == 1, 'Number of collisions in simple 2-body problem is wrong' # Next collision is at dt = 2sqrt(2) => dt=2.82 ==> t =~ 4.24 tx, tv, ncoll, _ = next(g) assert ncoll == 1, 'Number of collisions in simple 2-body problem is wrong' tx, tv, ncoll, _ = next(g) assert ncoll == 1, 'Number of collisions in simple 2-body problem is wrong' tx, tv, ncoll, _ = next(g) # collision should have happened now assert ncoll == 2, 'Number of collisions in simple 2-body problem is wrong' return None
def test_energy_conservation_sech2disk_manyparticles(): # Test that energy is conserved for a self-gravitating disk N = 101 totmass = 1. sigma = 1. zh = 2. * sigma**2. / totmass x = numpy.arctanh(2. * numpy.random.uniform(size=N) - 1) * zh v = numpy.random.normal(size=N) * sigma v -= numpy.mean(v) # stabilize m = numpy.ones_like(x) / N * (1. + 0.1 * (2. * numpy.random.uniform(size=N) - 1)) omega = 1.1 g = wendy.nbody(x, v, m, 0.05, omega=omega, approx=True, nleap=1000) E = wendy.energy(x, v, m, omega=omega) cnt = 0 while cnt < 100: tx, tv = next(g) assert numpy.fabs( wendy.energy(tx, tv, m, omega=omega) - E ) / E < 10.**-6., "Energy not conserved during approximate N-body integration with external harmonic potential" cnt += 1 return None
def test_notracermasses(): # approx should work with tracer sheets # Test that energy is conserved for a self-gravitating disk N = 101 totmass = 1. sigma = 1. zh = 2. * sigma**2. / totmass x = numpy.arctanh(2. * numpy.random.uniform(size=N) - 1) * zh v = numpy.random.normal(size=N) * sigma v -= numpy.mean(v) # stabilize m = numpy.ones_like(x) / N * (1. + 0.1 * (2. * numpy.random.uniform(size=N) - 1)) m[N // 2:] = 0. m *= 2. g = wendy.nbody(x, v, m, 0.05, approx=True, nleap=1000) E = wendy.energy(x, v, m) cnt = 0 while cnt < 100: tx, tv = next(g) assert numpy.fabs( wendy.energy(tx, tv, m) - E ) / E < 10.**-6., "Energy not conserved during approximate N-body integration with some tracer particles" cnt += 1 return None
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, fit_zsun=False,fit_omega=False, skipomega=10,delta_omega=0.3): """ 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 - potential parameter omega 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= (0.3) 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 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]],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) """ 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 w_prior is None: w_prior= w_out # 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]) Q_out= [] if output_wevolution: rndindx= numpy.random.permutation(len(w_out))[:output_wevolution] wevol= numpy.zeros((output_wevolution,npop,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) 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,use_v2=use_v2) 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) #ocounter= skipomega-1 # Causes F(omega) to be computed in the 1st step #z_prev, vz_prev= Aphi_to_zvz(A_init,phi_init-skipomega*step*omega_m2m, # omega_m2m) #Rewind for first step 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 omega (skipped in the first step, so undeclared vars okay) if fit_omega and ocounter == skipomega: domega= eps[1+fit_zsun]*step*skipomega*fco max_domega= delta_omega/30. 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 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= wendy.nbody(z_m2m[relevant_particles_index], vz_m2m[relevant_particles_index], mass[relevant_particles_index], step,maxcoll=10000000) 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) 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, z_obs,dens_obs_noise,tdelta_m2m, densv2_obs_noise,tdeltav2_m2m, kernel=kernel, kernel_deriv=kernel_deriv, h_m2m=h_m2m,use_v2=use_v2) 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 fco_new= force_of_change_omega(w_out,zsun_m2m,omega_m2m, z_m2m,vz_m2m,z_prev,vz_prev, step*skipomega, z_obs,dens_obs,dens_obs_noise, tdelta_m2m, densv2_obs,densv2_obs_noise, tdeltav2_m2m, h_m2m=h_m2m,kernel=kernel, delta_omega=delta_omega, use_v2=use_v2) 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(delta_m2m,delta_m2m_new)] fcw= fcw_new if fit_zsun: fcz= fcz_new if fit_omega and ocounter == skipomega: fco= fco_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_omega and ocounter == skipomega: fco+= step*skipomega*smooth*(fco_new-fco) else: fcw= fcw_new if fit_zsun: fcz= fcz_new if fit_omega and ocounter == skipomega: fco= fco_new # Record random weights if requested if output_wevolution: wevol[:,:,ii]= w_out[rndindx] out= (w_out,) if fit_zsun: out= out+(zsun_out,) if fit_omega: out= out+(omega_out,) out= out+(z_m2m,vz_m2m,) out= out+(numpy.array(Q_out),) if output_wevolution: out= out+(wevol,rndindx,) return out