1import numpy as np
2
3
4class FluidOperator:
5    def apply(self, ds):
6        for g in ds.index.grids:
7            self(g)
8
9
10class TopHatSphere(FluidOperator):
11    def __init__(self, radius, center, fields):
12        self.radius = radius
13        self.center = center
14        self.fields = fields
15
16    def __call__(self, grid, sub_select=None):
17        r = np.zeros(grid.ActiveDimensions, dtype="float64")
18        for i, ax in enumerate("xyz"):
19            np.add(r, (grid[ax].to_ndarray() - self.center[i]) ** 2.0, r)
20        np.sqrt(r, r)
21        ind = r <= self.radius
22        if sub_select is not None:
23            ind &= sub_select
24        for field, val in self.fields.items():
25            grid[field][ind] = val
26
27
28class CoredSphere(FluidOperator):
29    def __init__(self, core_radius, radius, center, fields):
30        self.radius = radius
31        self.center = center
32        self.fields = fields
33        self.core_radius = core_radius
34
35    def __call__(self, grid, sub_select=None):
36        r = np.zeros(grid.ActiveDimensions, dtype="float64")
37        r2 = self.radius ** 2
38        cr2 = self.core_radius ** 2
39        for i, ax in enumerate("xyz"):
40            np.add(r, (grid[ax] - self.center[i]) ** 2.0, r)
41        np.maximum(r, cr2, r)
42        ind = r <= r2
43        if sub_select is not None:
44            ind &= sub_select
45        for field, (outer_val, inner_val) in self.fields.items():
46            val = ((r[ind] - cr2) / (r2 - cr2)) ** 0.5 * (outer_val - inner_val)
47            grid[field][ind] = val + inner_val
48
49
50class BetaModelSphere(FluidOperator):
51    def __init__(self, beta, core_radius, radius, center, fields):
52        self.radius = radius
53        self.center = center
54        self.fields = fields
55        self.core_radius = core_radius
56        self.beta = beta
57
58    def __call__(self, grid, sub_select=None):
59        r = np.zeros(grid.ActiveDimensions, dtype="float64")
60        r2 = self.radius ** 2
61        cr2 = self.core_radius ** 2
62        for i, ax in enumerate("xyz"):
63            np.add(r, (grid[ax].ndarray_view() - self.center[i]) ** 2.0, r)
64        ind = r <= r2
65        if sub_select is not None:
66            ind &= sub_select
67        for field, core_val in self.fields.items():
68            val = core_val * (1.0 + r[ind] / cr2) ** (-1.5 * self.beta)
69            grid[field][ind] = val
70
71
72class NFWModelSphere(FluidOperator):
73    def __init__(self, scale_radius, radius, center, fields):
74        self.radius = radius
75        self.center = center
76        self.fields = fields
77        self.scale_radius = scale_radius  # unitless
78
79    def __call__(self, grid, sub_select=None):
80        r = np.zeros(grid.ActiveDimensions, dtype="float64")
81        for i, ax in enumerate("xyz"):
82            np.add(r, (grid[ax].ndarray_view() - self.center[i]) ** 2.0, r)
83        np.sqrt(r, r)
84        ind = r <= self.radius
85        r /= self.scale_radius
86        if sub_select is not None:
87            ind &= sub_select
88        for field, scale_val in self.fields.items():
89            val = scale_val / (r[ind] * (1.0 + r[ind]) ** 2)
90            grid[field][ind] = val
91
92
93class RandomFluctuation(FluidOperator):
94    def __init__(self, fields):
95        self.fields = fields
96
97    def __call__(self, grid, sub_select=None):
98        if sub_select is None:
99            sub_select = Ellipsis
100        for field, mag in self.fields.items():
101            vals = grid[field][sub_select]
102            rc = 1.0 + (np.random.random(vals.shape) - 0.5) * mag
103            grid[field][sub_select] *= rc
104