def smooth(t, inflection=10.0): error = sigmoid(-inflection / 2) return np.clip( (sigmoid(inflection * (t - 0.5)) - error) / (1 - 2 * error), 0, 1, )
class VectorField(VGroup): CONFIG = { "delta_x": 0.5, "delta_y": 0.5, "x_min": int(np.floor(-FRAME_WIDTH / 2)), "x_max": int(np.ceil(FRAME_WIDTH / 2)), "y_min": int(np.floor(-FRAME_HEIGHT / 2)), "y_max": int(np.ceil(FRAME_HEIGHT / 2)), "min_magnitude": 0, "max_magnitude": 2, "colors": DEFAULT_SCALAR_FIELD_COLORS, # Takes in actual norm, spits out displayed norm "length_func": lambda norm: 0.45 * sigmoid(norm), "opacity": 1.0, "vector_config": {}, } def __init__(self, func, **kwargs): VGroup.__init__(self, **kwargs) self.func = func self.rgb_gradient_function = get_rgb_gradient_function( self.min_magnitude, self.max_magnitude, self.colors, flip_alphas=False ) x_range = np.arange( self.x_min, self.x_max + self.delta_x, self.delta_x ) y_range = np.arange( self.y_min, self.y_max + self.delta_y, self.delta_y ) for x, y in it.product(x_range, y_range): point = x * RIGHT + y * UP self.add(self.get_vector(point)) self.set_opacity(self.opacity) def get_vector(self, point, **kwargs): output = np.array(self.func(point)) norm = get_norm(output) if norm == 0: output *= 0 else: output *= self.length_func(norm) / norm vector_config = dict(self.vector_config) vector_config.update(kwargs) vect = Vector(output, **vector_config) vect.shift(point) fill_color = rgb_to_color( self.rgb_gradient_function(np.array([norm]))[0] ) vect.set_color(fill_color) return vect
class VectorField(VGroup): CONFIG = { "step_multiple": 0.5, "magnitude_range": (0, 2), "color_map": "3b1b_colormap", # Takes in actual norm, spits out displayed norm "length_func": lambda norm: 0.45 * sigmoid(norm), "opacity": 1.0, "vector_config": {}, } def __init__( self, func: Callable[[float, float], Sequence[float]], coordinate_system: CoordinateSystem, **kwargs ): super().__init__(**kwargs) self.func = func self.coordinate_system = coordinate_system self.value_to_rgb = get_rgb_gradient_function( *self.magnitude_range, self.color_map, ) samples = get_sample_points_from_coordinate_system( coordinate_system, self.step_multiple ) self.add(*( self.get_vector(coords) for coords in samples )) def get_vector(self, coords: Iterable[float], **kwargs) -> Arrow: vector_config = merge_dicts_recursively( self.vector_config, kwargs ) output = np.array(self.func(*coords)) norm = get_norm(output) if norm > 0: output *= self.length_func(norm) / norm origin = self.coordinate_system.get_origin() _input = self.coordinate_system.c2p(*coords) _output = self.coordinate_system.c2p(*output) vect = Arrow( origin, _output, buff=0, **vector_config ) vect.shift(_input - origin) vect.set_rgba_array([[*self.value_to_rgb(norm), self.opacity]]) return vect