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