1 #include "sph/initial/Galaxy.h"
2 #include "gravity/BarnesHut.h"
3 #include "gravity/BruteForceGravity.h"
4 #include "gravity/Moments.h"
5 #include "math/rng/Rng.h"
6 #include "quantities/IMaterial.h"
7 #include "quantities/Quantity.h"
8 #include "quantities/Storage.h"
9 #include "system/Factory.h"
10 #include "system/Profiler.h"
11 #include "system/Settings.impl.h"
12 #include "system/Statistics.h"
13 
14 NAMESPACE_SPH_BEGIN
15 
16 //-----------------------------------------------------------------------------------------------------------
17 // Galaxy
18 //-----------------------------------------------------------------------------------------------------------
19 
20 // clang-format off
21 template<>
getDefaultSettings()22 const Settings<GalaxySettingsId>& getDefaultSettings() {
23     static Settings<GalaxySettingsId> instance({
24     { GalaxySettingsId::DISK_PARTICLE_COUNT,    "disk.particle_count",  10000, "" },
25     { GalaxySettingsId::DISK_RADIAL_CUTOFF,     "disk.radial_cutoff",   7.5_f, "" },
26     { GalaxySettingsId::DISK_RADIAL_SCALE,      "disk.radial_scale",    1._f, "" },
27     { GalaxySettingsId::DISK_VERTICAL_SCALE,    "disk.vertical_scale",  0.2_f, "" },
28     { GalaxySettingsId::DISK_VERTICAL_CUTOFF,   "disk.vertical_cutoff", 0.6_f, "" },
29     { GalaxySettingsId::DISK_TOOMRE_Q,          "disk.toomre_q",        1.2_f, "" },
30     { GalaxySettingsId::DISK_MASS,              "disk.mass",            1._f, "" },
31     { GalaxySettingsId::HALO_PARTICLE_COUNT,    "halo.particle_count",  10000, "" },
32     { GalaxySettingsId::HALO_SCALE_LENGTH,      "halo.scale_length",    10._f, "" },
33     { GalaxySettingsId::HALO_GAMMA,             "halo.gamma",           2._f, "" },
34     { GalaxySettingsId::HALO_CUTOFF,            "halo.cutoff",          15._f, "" },
35     { GalaxySettingsId::HALO_MASS,              "halo.mass",            5._f, "" },
36     { GalaxySettingsId::BULGE_PARTICLE_COUNT,   "bulge.particle_count", 10000, "" },
37     { GalaxySettingsId::BULGE_SCALE_LENGTH,     "bulge.scale_length",   0.4_f, "" },
38     { GalaxySettingsId::BULGE_CUTOFF,           "bulge.cutoff",         5._f, "" },
39     { GalaxySettingsId::BULGE_MASS,             "bulge.mass",           0.6_f, "" },
40     { GalaxySettingsId::PARTICLE_RADIUS,        "particle_radius",      0.01_f, "" },
41     });
42     return instance;
43 }
44 // clang-format on
45 
46 // Explicit instantiation
47 template class Settings<GalaxySettingsId>;
48 template class SettingsIterator<GalaxySettingsId>;
49 
50 /// Mostly uses methods described in https://github.com/nmuldavin/NBodyIntegrator
51 
52 /// Surface probability distribution of a disk
diskSurfacePdf(const Float r,const Float h)53 INLINE Float diskSurfacePdf(const Float r, const Float h) {
54     return exp(-r / h) * r;
55 }
56 
57 /// Normalized surface density of a disk
diskSurfaceDensity(const Float r,const Float h,const Float m_disk)58 INLINE Float diskSurfaceDensity(const Float r, const Float h, const Float m_disk) {
59     return m_disk / (2._f * PI * sqr(h)) * exp(-r / h);
60 }
61 
62 /// Vertical mass distribution of a disk
diskVerticalPdf(const Float z,const Float z0)63 INLINE Float diskVerticalPdf(const Float z, const Float z0) {
64     return 1._f / sqr(cosh(z / z0));
65 }
66 
67 /// Probability distribution function of a halo
haloPdf(const Float r,const Float r0,const Float g0)68 INLINE Float haloPdf(const Float r, const Float r0, const Float g0) {
69     return exp(-sqr(r / r0)) / (sqr(r) + sqr(g0)) * sqr(r);
70 }
71 
maxHaloPdf(const Float r0,const Float g0)72 INLINE Float maxHaloPdf(const Float r0, const Float g0) {
73     const Float x2 = 0.5_f * (sqrt(sqr(g0) * (sqr(g0) + 4._f * sqr(r0))) - sqr(g0));
74     SPH_ASSERT(x2 > 0._f);
75     return haloPdf(sqrt(x2), r0, g0);
76 }
77 
78 /// Probability distribution function for velocities in halo and bulge.
velocityPdf(const Float v,const Float sigma_r)79 INLINE Float velocityPdf(const Float v, const Float sigma_r) {
80     return sqr(v) * exp(-0.5_f * sqr(v) / sigma_r);
81 }
82 
83 /// Probability distribution function of a bulge
bulgePdf(const Float r,const Float a)84 INLINE Float bulgePdf(const Float r, const Float a) {
85     return r / (sqr(a) * pow<3>(1._f + r / a));
86 }
87 
getEpicyclicFrequency(IGravity & gravity,const Vector & r,const Vector & dv1,const Float dr)88 static Float getEpicyclicFrequency(IGravity& gravity, const Vector& r, const Vector& dv1, const Float dr) {
89     const Float radius = sqrt(sqr(r[X]) + sqr(r[Y])) + EPS;
90     const Vector dv2 = gravity.evalAcceleration(r * (1._f + dr));
91 
92     const Float a1_rad = (dv1[X] * r[X] + dv1[Y] * r[Y]) / radius;
93     const Float a2_rad = (dv2[X] * r[X] + dv2[Y] * r[Y]) / radius;
94 
95     const Float k2 = (3._f / radius) * a1_rad + (a2_rad - a1_rad) / dr;
96     return sqrt(abs(k2));
97 }
98 
generateDisk(UniformRng & rng,const GalaxySettings & settings)99 Storage Galaxy::generateDisk(UniformRng& rng, const GalaxySettings& settings) {
100     MEASURE_SCOPE("Galaxy::generateDisk");
101 
102     Array<Vector> positions;
103     const Size n_disk = settings.get<int>(GalaxySettingsId::DISK_PARTICLE_COUNT);
104     const Float r_cutoff = settings.get<Float>(GalaxySettingsId::DISK_RADIAL_CUTOFF);
105     const Float r0 = settings.get<Float>(GalaxySettingsId::DISK_RADIAL_SCALE);
106     const Float z_cutoff = settings.get<Float>(GalaxySettingsId::DISK_VERTICAL_CUTOFF);
107     const Float z0 = settings.get<Float>(GalaxySettingsId::DISK_VERTICAL_SCALE);
108     const Float h = settings.get<Float>(GalaxySettingsId::PARTICLE_RADIUS);
109 
110     const Interval radialRange(0._f, r_cutoff);
111     const Interval verticalRange(-z_cutoff, z_cutoff);
112 
113     // radial PDF is maximal at r = r0
114     const Float maxSurfacePdf = diskSurfacePdf(r0, r0);
115 
116     // vertical pdf is maximal at z = 0
117     const Float maxVerticalPdf = diskVerticalPdf(0, z0);
118 
119     for (Size i = 0; i < n_disk; ++i) {
120         const Float r = sampleDistribution(
121             rng, radialRange, maxSurfacePdf, [r0](const Float x) { return diskSurfacePdf(x, r0); });
122 
123         const Float phi = rng() * 2._f * PI;
124 
125         const Float z = sampleDistribution(
126             rng, verticalRange, maxVerticalPdf, [z0](const Float x) { return diskVerticalPdf(x, z0); });
127 
128         Vector pos = cylindricalToCartesian(r, phi, z);
129         pos[H] = h;
130         positions.push(pos);
131     }
132 
133 
134     const Float m_disk = settings.get<Float>(GalaxySettingsId::DISK_MASS);
135     const Float m = m_disk / n_disk;
136 
137     Storage storage(makeShared<NullMaterial>(BodySettings::getDefaults()));
138     storage.insert<Vector>(QuantityId::POSITION, OrderEnum::SECOND, std::move(positions));
139     storage.insert<Float>(QuantityId::MASS, OrderEnum::ZERO, m);
140     storage.insert<Size>(QuantityId::FLAG, OrderEnum::ZERO, Size(PartEnum::DISK));
141     return storage;
142 }
143 
generateHalo(UniformRng & rng,const GalaxySettings & settings)144 Storage Galaxy::generateHalo(UniformRng& rng, const GalaxySettings& settings) {
145     MEASURE_SCOPE("Galaxy::generateHalo");
146 
147     const Size n_halo = settings.get<int>(GalaxySettingsId::HALO_PARTICLE_COUNT);
148     const Float cutoff = settings.get<Float>(GalaxySettingsId::HALO_CUTOFF);
149     const Float r0 = settings.get<Float>(GalaxySettingsId::HALO_SCALE_LENGTH);
150     const Float g0 = settings.get<Float>(GalaxySettingsId::HALO_GAMMA);
151     const Float h = settings.get<Float>(GalaxySettingsId::PARTICLE_RADIUS);
152     const Interval range(0._f, cutoff);
153 
154     const Float maxPdf = maxHaloPdf(r0, g0);
155 
156     Array<Vector> positions;
157     for (Size i = 0; i < n_halo; ++i) {
158         const Float r = sampleDistribution(rng, range, maxPdf, [r0, g0](const Float x) { //
159             return haloPdf(x, r0, g0);
160         });
161 
162         Vector pos = sampleUnitSphere(rng) * r;
163         pos[H] = h;
164         positions.push(pos);
165     }
166 
167     const Float m_halo = settings.get<Float>(GalaxySettingsId::HALO_MASS);
168     const Float m = m_halo / n_halo;
169 
170     Storage storage(makeShared<NullMaterial>(BodySettings::getDefaults()));
171     storage.insert<Vector>(QuantityId::POSITION, OrderEnum::SECOND, std::move(positions));
172     storage.insert<Float>(QuantityId::MASS, OrderEnum::ZERO, m);
173     storage.insert<Size>(QuantityId::FLAG, OrderEnum::ZERO, Size(PartEnum::HALO));
174     return storage;
175 }
176 
generateBulge(UniformRng & rng,const GalaxySettings & settings)177 Storage Galaxy::generateBulge(UniformRng& rng, const GalaxySettings& settings) {
178     MEASURE_SCOPE("Galaxy::generateBulge");
179 
180     const Size n_bulge = settings.get<int>(GalaxySettingsId::HALO_PARTICLE_COUNT);
181     const Float cutoff = settings.get<Float>(GalaxySettingsId::BULGE_CUTOFF);
182     const Float a = settings.get<Float>(GalaxySettingsId::BULGE_SCALE_LENGTH);
183     const Float h = settings.get<Float>(GalaxySettingsId::PARTICLE_RADIUS);
184     const Interval range(0._f, cutoff);
185 
186     // PDF is maximal at x=a/2
187     const Float maxPdf = bulgePdf(0.5_f * a, a);
188 
189     Array<Vector> positions;
190     for (Size i = 0; i < n_bulge; ++i) {
191         const Float r = sampleDistribution(rng, range, maxPdf, [a](const Float x) { return bulgePdf(x, a); });
192 
193         Vector pos = sampleUnitSphere(rng) * r;
194         pos[H] = h;
195         positions.push(pos);
196     }
197 
198     const Float m_bulge = settings.get<Float>(GalaxySettingsId::BULGE_MASS);
199     const Float m = m_bulge / n_bulge;
200 
201     Storage storage(makeShared<NullMaterial>(BodySettings::getDefaults()));
202     storage.insert<Vector>(QuantityId::POSITION, OrderEnum::SECOND, std::move(positions));
203     storage.insert<Float>(QuantityId::MASS, OrderEnum::ZERO, m);
204     storage.insert<Size>(QuantityId::FLAG, OrderEnum::ZERO, Size(PartEnum::BULGE));
205     return storage;
206 }
207 
computeCumulativeMass(const GalaxySettings & settings,const Storage & storage)208 static Array<Pair<Float>> computeCumulativeMass(const GalaxySettings& settings, const Storage& storage) {
209     MEASURE_SCOPE("computeCumulativeMass");
210 
211     constexpr Size MASS_BINS = 1000;
212 
213     const Float haloCutoff = settings.get<Float>(GalaxySettingsId::HALO_CUTOFF);
214     const Float dr = haloCutoff / MASS_BINS;
215 
216     ArrayView<const Vector> r = storage.getValue<Vector>(QuantityId::POSITION);
217     ArrayView<const Float> m = storage.getValue<Float>(QuantityId::MASS);
218 
219     Array<Pair<Float>> cumulativeDist;
220     Array<Float> differentialDist(MASS_BINS);
221     differentialDist.fill(0._f);
222     for (Size i = 0; i < r.size(); ++i) {
223         const Float radius = getLength(r[i]);
224         const Size binIdx = Size(radius * MASS_BINS / haloCutoff);
225         SPH_ASSERT(binIdx < MASS_BINS);
226 
227         differentialDist[min(binIdx, MASS_BINS - 1)] += m[i];
228     }
229 
230     Float massSum = 0._f;
231     for (Size binIdx = 0; binIdx < MASS_BINS; ++binIdx) {
232         const Float binRadius = (binIdx + 1) * dr;
233         massSum += differentialDist[binIdx];
234         cumulativeDist.push(Pair<Float>{ binRadius, massSum });
235     }
236 
237     return cumulativeDist;
238 }
239 
getPartSequence(const Storage & storage,const Galaxy::PartEnum id)240 static IndexSequence getPartSequence(const Storage& storage, const Galaxy::PartEnum id) {
241     ArrayView<const Size> flag = storage.getValue<Size>(QuantityId::FLAG);
242     Iterator<const Size> iterFrom, iterTo;
243     std::tie(iterFrom, iterTo) = std::equal_range(flag.begin(), flag.end(), Size(id));
244     return IndexSequence(
245         Size(std::distance(flag.begin(), iterFrom)), Size(std::distance(flag.begin(), iterTo)));
246 }
247 
computeDiskVelocities(IScheduler & scheduler,UniformRng & rng,const GalaxySettings & settings,Storage & storage)248 static void computeDiskVelocities(IScheduler& scheduler,
249     UniformRng& rng,
250     const GalaxySettings& settings,
251     Storage& storage) {
252     MEASURE_SCOPE("computeDiskVelocities");
253 
254     const Float r0 = settings.get<Float>(GalaxySettingsId::DISK_RADIAL_SCALE);
255     const Float z0 = settings.get<Float>(GalaxySettingsId::DISK_VERTICAL_SCALE);
256     const Float r_ref = 2.5_f * r0;
257     const Float r_cutoff = settings.get<Float>(GalaxySettingsId::DISK_RADIAL_CUTOFF);
258     const Float m_disk = settings.get<Float>(GalaxySettingsId::DISK_MASS);
259     const Float Q = settings.get<Float>(GalaxySettingsId::DISK_TOOMRE_Q);
260     // const Float double as = 0.25 * h;
261     const Float dr = 1.e-3_f * r_cutoff;
262     const Float as = 0.25_f * r0;
263 
264     ArrayView<Vector> r, v, dv;
265     tie(r, v, dv) = storage.getAll<Vector>(QuantityId::POSITION);
266 
267     const IndexSequence sequence = getPartSequence(storage, Galaxy::PartEnum::DISK);
268 
269     BarnesHut gravity(0.8_f, MultipoleOrder::OCTUPOLE, SolidSphereKernel{}, 25, 50, 1._f);
270     // BruteForceGravity gravity(SolidSphereKernel{}, 1._f);
271     gravity.build(scheduler, storage);
272     Statistics stats;
273     std::fill(dv.begin(), dv.end(), Vector(0._f));
274     gravity.evalSelfGravity(scheduler, dv, stats);
275 
276     Float sigma = 0._f;
277     Size count = 0;
278     Float annulus = dr;
279     while (count == 0._f) {
280         for (Size i : sequence) {
281             const Float radius = sqrt(sqr(r[i][X]) + sqr(r[i][Y]));
282             if (abs(radius - r_ref) < annulus) {
283                 const Float kappa = getEpicyclicFrequency(gravity, r[i], dv[i], 0.05_f * annulus);
284                 sigma += 3.36_f * diskSurfaceDensity(radius, r0, m_disk) / kappa;
285                 count++;
286             }
287         }
288         // if no particle lies in the annulus, try again with larger value
289         annulus *= 2._f;
290     }
291     SPH_ASSERT(count > 0);
292 
293     sigma = sigma * Q / count;
294 
295     const Float A = sqr(sigma) / diskSurfaceDensity(r_ref, r0, m_disk);
296     SPH_ASSERT(A >= 0._f, A);
297 
298     parallelFor(scheduler, sequence, [&](const Size i) {
299         const Float radius = sqrt(sqr(r[i][X]) + sqr(r[i][Y]));
300         const Float vz2 = PI * z0 * diskSurfaceDensity(sqrt(sqr(radius) + 2._f * sqr(as)), r0, m_disk);
301         const Float vz = sampleNormalDistribution(rng, 0._f, vz2);
302         SPH_ASSERT(vz2 > 0._f);
303 
304         const Float vr2 = A * vz2 / (PI * z0);
305         const Float vr = sampleNormalDistribution(rng, 0._f, vr2);
306         SPH_ASSERT(vr2 > 0._f);
307 
308         const Vector a = dv[i];
309         const Float ar = (a[X] * r[i][X] + a[Y] * r[i][Y]) / radius;
310         SPH_ASSERT(isReal(ar));
311 
312         const Float omega = sqrt(abs(ar) / radius);
313         SPH_ASSERT(isReal(omega));
314 
315         const Float kappa = getEpicyclicFrequency(gravity, r[i], dv[i], dr);
316         SPH_ASSERT(isReal(kappa));
317 
318         //  circular velocity
319         const Float v_c = omega * radius;
320         Float va = sqrt(abs(sqr(v_c) + vr2 * (1._f - sqr(kappa) / (4._f * sqr(omega)) - 2._f * radius / r0)));
321         SPH_ASSERT(isReal(va));
322 
323         const Float sigma2 = vr2 * sqr(kappa) / (4._f * sqr(omega));
324         va += sampleNormalDistribution(rng, 0._f, sigma2);
325 
326         //  transform to cartesian coordinates
327         const Float c = r[i][X] / radius;
328         const Float s = r[i][Y] / radius;
329         v[i][X] = vr * c - va * s;
330         v[i][Y] = vr * s + va * c;
331         v[i][Z] = vz;
332     });
333 }
334 
335 template <typename TFunc>
computeSphericalVelocities(UniformRng & rng,ArrayView<const Pair<Float>> massDist,const Galaxy::PartEnum partId,Storage & storage,const TFunc & func)336 static void computeSphericalVelocities(UniformRng& rng,
337     ArrayView<const Pair<Float>> massDist,
338     const Galaxy::PartEnum partId,
339     Storage& storage,
340     const TFunc& func) {
341     const Float dr = massDist[1][0] - massDist[0][0];
342 
343     ArrayView<const Vector> r = storage.getValue<Vector>(QuantityId::POSITION);
344     ArrayView<Vector> v = storage.getDt<Vector>(QuantityId::POSITION);
345 
346     const IndexSequence sequence = getPartSequence(storage, partId);
347     for (Size i : sequence) {
348         const Float radius = getLength(r[i]);
349         const Size firstBin = Size(radius / dr);
350 
351         const Float v_esc = sqrt(2._f * massDist[firstBin][1] / radius);
352 
353         Float vr2 = 0._f;
354         for (Size binIdx = firstBin; binIdx < massDist.size(); ++binIdx) {
355             vr2 += func(massDist[binIdx][0]) * dr * massDist[binIdx][1];
356         }
357         vr2 /= func(radius) / sqr(radius);
358 
359         const Interval range(0._f, 0.95_f * v_esc);
360         const Float maxPdf = velocityPdf(sqrt(2._f * vr2), vr2);
361 
362         const Float u = sampleDistribution(rng, range, maxPdf, [vr2](const Float x) { //
363             return velocityPdf(x, vr2);
364         });
365 
366         v[i] = sampleUnitSphere(rng) * u;
367     }
368 }
369 
computeHaloVelocities(UniformRng & rng,const GalaxySettings & settings,ArrayView<const Pair<Float>> massDist,Storage & storage)370 static void computeHaloVelocities(UniformRng& rng,
371     const GalaxySettings& settings,
372     ArrayView<const Pair<Float>> massDist,
373     Storage& storage) {
374     MEASURE_SCOPE("computeHaloVelocities");
375 
376     const Float r0 = settings.get<Float>(GalaxySettingsId::HALO_SCALE_LENGTH);
377     const Float g0 = settings.get<Float>(GalaxySettingsId::HALO_GAMMA);
378 
379     computeSphericalVelocities(rng, massDist, Galaxy::PartEnum::HALO, storage, [r0, g0](const Float x) { //
380         return haloPdf(x, r0, g0);
381     });
382 }
383 
computeBulgeVelocities(UniformRng & rng,const GalaxySettings & settings,ArrayView<const Pair<Float>> massDist,Storage & storage)384 static void computeBulgeVelocities(UniformRng& rng,
385     const GalaxySettings& settings,
386     ArrayView<const Pair<Float>> massDist,
387     Storage& storage) {
388     MEASURE_SCOPE("computeBulgeVelocities");
389 
390     const Float a = settings.get<Float>(GalaxySettingsId::BULGE_SCALE_LENGTH);
391 
392     computeSphericalVelocities(rng, massDist, Galaxy::PartEnum::BULGE, storage, [a](const Float x) { //
393         return bulgePdf(x, a);
394     });
395 }
396 
397 class StorageBuilder {
398 public:
399     Storage storage;
400     const Galaxy::IProgressCallbacks& callbacks;
401     Size partId = 0;
402 
403     const Size numParts = 8;
404 
405 public:
StorageBuilder(const Galaxy::IProgressCallbacks & callbacks)406     StorageBuilder(const Galaxy::IProgressCallbacks& callbacks)
407         : callbacks(callbacks) {}
408 
operator *()409     Storage& operator*() {
410         callbacks.onPart(storage, partId, numParts);
411         ++partId;
412 
413         return storage;
414     }
415 
operator ->()416     Storage* operator->() {
417         return &operator*();
418     }
419 
release()420     Storage release() && {
421         return std::move(storage);
422     }
423 };
424 
generateIc(const RunSettings & globals,const GalaxySettings & settings,const IProgressCallbacks & callbacks)425 Storage Galaxy::generateIc(const RunSettings& globals,
426     const GalaxySettings& settings,
427     const IProgressCallbacks& callbacks) {
428     const int seed = globals.get<int>(RunSettingsId::RUN_RNG_SEED);
429     UniformRng rng(seed);
430     SharedPtr<IScheduler> scheduler = Factory::getScheduler(globals);
431 
432     StorageBuilder builder(callbacks);
433     builder->merge(generateDisk(rng, settings));
434     builder->merge(generateHalo(rng, settings));
435     builder->merge(generateBulge(rng, settings));
436 
437     Array<Pair<Float>> massDist = computeCumulativeMass(settings, *builder);
438     computeDiskVelocities(*scheduler, rng, settings, *builder);
439     computeHaloVelocities(rng, settings, massDist, *builder);
440     computeBulgeVelocities(rng, settings, massDist, *builder);
441 
442     Storage storage = std::move(builder).release();
443     ArrayView<const Size> flag = storage.getValue<Size>(QuantityId::FLAG);
444     SPH_ASSERT(std::is_sorted(flag.begin(), flag.end()));
445     return storage;
446 }
447 
448 NAMESPACE_SPH_END
449