""" Code from Modeling and Simulation in Python. Copyright 2017 Allen Downey License: https://creativecommons.org/licenses/by/4.0) """ import logging logger = logging.getLogger(name='modsim.py') #TODO: Make this Python 3.7 when conda is ready # make sure we have Python 3.6 or better import sys if sys.version_info < (3, 6): logger.warning('modsim.py depends on Python 3.6 features.') import inspect import matplotlib.pyplot as plt import numpy as np import pandas as pd import scipy import sympy import seaborn as sns sns.set(style='white', font_scale=1.2) import pint UNITS = pint.UnitRegistry() Quantity = UNITS.Quantity # expose some names so we can use them without dot notation from copy import copy from numpy import sqrt, log, exp, pi from pandas import DataFrame, Series from time import sleep from scipy.interpolate import interp1d from scipy.interpolate import InterpolatedUnivariateSpline from scipy.integrate import odeint from scipy.integrate import solve_ivp from scipy.optimize import leastsq from scipy.optimize import minimize_scalar import scipy.optimize def flip(p=0.5): """Flips a coin with the given probability. p: float 0-1 returns: boolean (True or False) """ return np.random.random() < p # For all the built-in Python functions that do math, # let's use the NumPy version instead. abs = np.abs min = np.min max = np.max pow = np.power sum = np.sum round = np.round def cart2pol(x, y, z=None): """Convert Cartesian coordinates to polar. x: number or sequence y: number or sequence z: number or sequence (optional) returns: theta, rho OR theta, rho, z """ x = np.asarray(x) y = np.asarray(y) # TODO: use hypot? rho = np.sqrt(x**2 + y**2) theta = np.arctan2(y, x) if z is None: return theta, rho else: return theta, rho, z def pol2cart(theta, rho, z=None): """Convert polar coordinates to Cartesian. theta: number or sequence rho: number or sequence z: number or sequence (optional) returns: x, y OR x, y, z """ if hasattr(theta, 'units'): if theta.units == UNITS.degree: theta = theta.to(UNITS.radian) if theta.units != UNITS.radian: msg = """In pol2cart, theta must be either a number or a Quantity in degrees or radians.""" raise ValueError(msg) x = rho * np.cos(theta) y = rho * np.sin(theta) if z is None: return x, y else: return x, y, z def linspace(start, stop, num=50, **options): """Returns an array of evenly-spaced values in the interval [start, stop]. start: first value stop: last value num: number of values Also accepts the same keyword arguments as np.linspace. See https://docs.scipy.org/doc/numpy/reference/generated/numpy.linspace.html returns: array or Quantity """ underride(options, dtype=np.float64) # see if either of the arguments has units units = getattr(start, 'units', None) units = getattr(stop, 'units', units) array = np.linspace(start, stop, num, **options) if units: array = array * units return array def linrange(start=0, stop=None, step=1, **options): """Returns an array of evenly-spaced values in the interval [start, stop]. This function works best if the space between start and stop is divisible by step; otherwise the results might be surprising. By default, the last value in the array is `stop-step` (at least approximately). If you provide the keyword argument `endpoint=True`, the last value in the array is `stop`. start: first value stop: last value step: space between values returns: array or Quantity """ if stop is None: stop = start start = 0 # TODO: what breaks if we don't make the dtype float? #underride(options, endpoint=True, dtype=np.float64) underride(options, endpoint=False) # see if any of the arguments has units units = getattr(start, 'units', None) units = getattr(stop, 'units', units) units = getattr(step, 'units', units) n = np.round((stop - start) / step) if options['endpoint']: n += 1 array = np.full(int(n), magnitude(step)) array[0] = magnitude(start) array = np.cumsum(array) if units: array = array * units return array def magnitude(x): """Returns the magnitude of a Quantity or number. x: Quantity or number returns: number """ return x.magnitude if isinstance(x, Quantity) else x def magnitudes(x): """Returns the magnitude of a Quantity or number, or sequence. x: Quantity or number, or sequence returns: number """ try: return [magnitude(elt) for elt in x] except TypeError: # not iterable return magnitude(x) def units(x): """Returns the units of a Quantity or number. x: Quantity or number returns: Unit object or 1 """ return x.units if isinstance(x, Quantity) else 1 def remove_units(series): """Removes units from the values in a Series. Only removes units from top-level values; does not traverse nested values. returns: new Series object """ res = copy(series) print(type(res)) for label, value in res.iteritems(): res[label] = magnitude(value) return res def require_units(x, units): """Apply units to `x`, if necessary. x: Quantity or number units: Pint Units object returns: Quantity """ if isinstance(x, Quantity): return x.to(units) else: return Quantity(x, units) def fit_leastsq(error_func, params, *args, **options): """Find the parameters that yield the best fit for the data. `params` can be a sequence, array, or Series Whatever arguments are provided are passed along to `error_func` error_func: function that computes a sequence of errors params: initial guess for the best parameters data: the data to be fit; will be passed to min_fun options: any other arguments are passed to leastsq """ # if any of the params are quantities, strip the units x0 = [magnitude(x) for x in params] # override `full_output` so we get a message if something goes wrong options['full_output'] = True # run leastsq with units_off(): best_params, cov_x, infodict, mesg, ier = leastsq(error_func, x0=x0, args=args, **options) details = ModSimSeries(infodict) details.set(cov_x=cov_x, mesg=mesg, ier=ier) # if we got a Params object, we should return a Params object if isinstance(params, Params): best_params = Params(Series(best_params, params.index)) # return the best parameters and details return best_params, details def min_bounded(min_func, bounds, *args, **options): """Finds the input value that minimizes `min_func`. Wrapper for https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize_scalar.html min_func: computes the function to be minimized bounds: sequence of two values, lower and upper bounds of the range to be searched args: any additional positional arguments are passed to min_func options: any keyword arguments are passed as options to minimize_scalar returns: ModSimSeries object """ # try: # print(bounds[0]) # min_func(bounds[0], *args) # except Exception as e: # msg = """Before running scipy.integrate.min_bounded, I tried # running the slope function you provided with the # initial conditions in system and t=0, and I got # the following error:""" # logger.error(msg) # raise(e) underride(options, xatol=1e-3) # TODO: Do we need to remove units from bounds? with units_off(): res = minimize_scalar(min_func, bracket=bounds, bounds=bounds, args=args, method='bounded', options=options) if not res.success: msg = """scipy.optimize.minimize_scalar did not succeed. The message it returned is %s""" % res.message raise Exception(msg) return ModSimSeries(res) def max_bounded(max_func, bounds, *args, **options): """Finds the input value that maximizes `max_func`. Wrapper for https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize_scalar.html min_func: computes the function to be maximized bounds: sequence of two values, lower and upper bounds of the range to be searched args: any additional positional arguments are passed to max_func options: any keyword arguments are passed as options to minimize_scalar returns: ModSimSeries object """ def min_func(*args): return -max_func(*args) res = min_bounded(min_func, bounds, *args, **options) # we have to negate the function value before returning res res.fun = -res.fun return res def minimize(min_func, x0, *args, **options): """Finds the input value that minimizes `min_func`. Wrapper for https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize.html min_func: computes the function to be minimized x0: initial guess args: any additional positional arguments are passed to min_func options: any keyword arguments are passed as options to minimize_scalar returns: ModSimSeries object """ underride(options, tol=1e-3) with units_off(): res = scipy.optimize.minimize(min_func, x0, *args, **options) return ModSimSeries(res) def run_odeint(system, slope_func, **options): """Integrates an ordinary differential equation. `system` should contain system parameters and `ts`, which is an array or Series that specifies the time when the solution will be computed. system: System object slope_func: function that computes slopes returns: TimeFrame """ # make sure `system` contains `ts` if not hasattr(system, 'ts'): msg = """It looks like `system` does not contain `ts` as a system variable. `ts` should be an array or Series that specifies the times when the solution will be computed:""" raise ValueError(msg) # make sure `system` contains `init` if not hasattr(system, 'init'): msg = """It looks like `system` does not contain `init` as a system variable. `init` should be a State object that specifies the initial condition:""" raise ValueError(msg) # make the system parameters available as globals unpack(system) # try running the slope function with the initial conditions try: slope_func(init, ts[0], system) except Exception as e: msg = """Before running scipy.integrate.odeint, I tried running the slope function you provided with the initial conditions in system and t=0, and I got the following error:""" logger.error(msg) raise(e) # when odeint calls slope_func, it should pass `system` as # the third argument. To make that work, we have to make a # tuple with a single element and pass the tuple to odeint as `args` args = (system,) # now we're ready to run `odeint` with `init` and `ts` from `system` with units_off(): array = odeint(slope_func, list(init), ts, args, **options) # the return value from odeint is an array, so let's pack it into # a TimeFrame with appropriate columns and index frame = TimeFrame(array, columns=init.index, index=ts, dtype=np.float64) return frame def run_ode_solver(system, slope_func, **options): """Computes a numerical solution to a differential equation. `system` must contain `init` with initial conditions, `t_0` with the start time, and `t_end` with the end time. It can contain any other parameters required by the slope function. `options` can be any legal options of `scipy.integrate.solve_ivp` system: System object slope_func: function that computes slopes returns: TimeFrame """ # make sure `system` contains `init` if not hasattr(system, 'init'): msg = """It looks like `system` does not contain `init` as a system variable. `init` should be a State object that specifies the initial condition:""" raise ValueError(msg) # make sure `system` contains `t_end` if not hasattr(system, 't_end'): msg = """It looks like `system` does not contain `t_end` as a system variable. `t_end` should be the final time:""" raise ValueError(msg) # make the system parameters available as globals unpack(system) # the default value for t_0 is 0 t_0 = getattr(system, 't_0', 0) # try running the slope function with the initial conditions # try: # slope_func(init, t_0, system) # except Exception as e: # msg = """Before running scipy.integrate.solve_ivp, I tried # running the slope function you provided with the # initial conditions in `system` and `t=t_0` and I got # the following error:""" # logger.error(msg) # raise(e) # wrap the slope function to reverse the arguments and add `system` f = lambda t, y: slope_func(y, t, system) def wrap_event(event): """Wrap the event functions. Make events terminal by default. """ wrapped = lambda t, y: event(y, t, system) wrapped.terminal = getattr(event, 'terminal', True) wrapped.direction = getattr(event, 'direction', 0) return wrapped # wrap the event functions so they take the right arguments events = options.pop('events', []) try: events = [wrap_event(event) for event in events] except TypeError: events = wrap_event(events) # remove dimensions from the initial conditions. # we need this because otherwise `init` gets copied into the # results array along with its units # try: # y_0 = [magnitude(x) for x in init] # except TypeError: # y_0 = [magnitude(init)] y_0 = [magnitude(x) for x in init] # run the solver with units_off(): bunch = solve_ivp(f, [t_0, t_end], y_0, events=events, **options) # separate the results from the details y = bunch.pop('y') t = bunch.pop('t') details = ModSimSeries(bunch) # pack the results into a TimeFrame results = TimeFrame(np.transpose(y), index=t, columns=init.index) return results, details def fsolve(func, x0, *args, **options): """Return the roots of the (non-linear) equations defined by func(x) = 0 given a starting estimate. Uses scipy.optimize.fsolve, with extra error-checking. func: function to find the roots of x0: scalar or array, initial guess args: additional positional arguments are passed along to fsolve, which passes them along to func returns: solution as an array """ # make sure we can run the given function with x0 try: func(x0, *args) except Exception as e: msg = """Before running scipy.optimize.fsolve, I tried running the error function you provided with the x0 you provided, and I got the following error:""" logger.error(msg) raise(e) # make the tolerance more forgiving than the default underride(options, xtol=1e-6) x0 = magnitude(x0) # run fsolve with units_off(): result = scipy.optimize.fsolve(func, x0, args=args, **options) return result def crossings(series, value): """Find the labels where the series passes through value. The labels in series must be increasing numerical values. series: Series value: number returns: sequence of labels """ interp = InterpolatedUnivariateSpline(series.index, series-value) return interp.roots() def interpolate(series, **options): """Creates an interpolation function. series: Series object options: any legal options to scipy.interpolate.interp1d returns: function that maps from the index of the series to values """ # TODO: add error checking for nonmonotonicity if sum(series.index.isnull()): msg = """The Series you passed to interpolate contains NaN values in the index, which would result in undefined behavior. So I'm putting a stop to that.""" raise ValueError(msg) # make the interpolate function extrapolate past the ends of # the range, unless `options` already specifies a value for `fill_value` underride(options, fill_value='extrapolate') # call interp1d, which returns a new function object interp_func = interp1d(series.index, series.values, **options) units = getattr(series, 'units', None) if units: return lambda x: Quantity(interp_func(x), units) else: return interp_func def interp_inverse(series, **options): """Interpolate the inverse function of a Series. series: Series object, represents a mapping from `a` to `b` kind: string, which kind of iterpolation options: keyword arguments passed to interpolate returns: interpolation object, can be used as a function from `b` to `a` """ inverse = Series(series.index, index=series.values) T = interpolate(inverse, **options) return T def unpack(series): """Make the names in `series` available as globals. series: Series with variables names in the index """ # TODO: Make this a context manager, so the syntax is # with series: # and maybe even add an __exit__ that copies changes back frame = inspect.currentframe() caller = frame.f_back caller.f_globals.update(series) def source_code(obj): """Prints the source code for a given object. obj: function or method object """ print(inspect.getsource(obj)) def underride(d, **options): """Add key-value pairs to d only if key is not in d. If d is None, create a new dictionary. d: dictionary options: keyword args to add to d """ if d is None: d = {} for key, val in options.items(): d.setdefault(key, val) return d def plot(*args, **options): """Makes line plots. args can be: plot(y) plot(y, style_string) plot(x, y) plot(x, y, style_string) options are the same as for pyplot.plot """ # TODO: add lines to REPLOT_CACHE x, y, style = parse_plot_args(*args, **options) if isinstance(x, pd.DataFrame) or isinstance(y, pd.DataFrame): raise ValueError("modsimpy.plot can't handle DataFrames.") if x is None: if isinstance(y, np.ndarray): x = np.arange(len(y)) if isinstance(y, pd.Series): x = y.index y = y.values x = magnitudes(x) y = magnitudes(y) underride(options, linewidth=3, alpha=0.6) if style is not None: lines = plt.plot(x, y, style, **options) else: lines = plt.plot(x, y, **options) return lines REPLOT_CACHE = {} def replot(*args, **options): """ """ try: label = options['label'] except KeyError: raise ValueError('To use replot, you must provide a label argument.') axes = plt.gca() key = (axes, label) if key not in REPLOT_CACHE: lines = plot(*args, **options) if len(lines) != 1: raise ValueError('Replot only works with a single plotted element.') REPLOT_CACHE[key] = lines[0] return lines line = REPLOT_CACHE[key] x, y, style = parse_plot_args(*args, **options) line.set_xdata(x) line.set_ydata(y) def parse_plot_args(*args, **options): """Parse the args the same way plt.plot does.""" x = None y = None style = None if len(args) == 1: y = args[0] elif len(args) == 2: if isinstance(args[1], str): y, style = args else: x, y = args elif len(args) == 3: x, y, style = args return x, y, style def contour(df, **options): """Makes a contour plot from a DataFrame. Note: columns and index must be numerical df: DataFrame """ x = df.columns y = df.index X, Y = np.meshgrid(x, y) cs = plt.contour(X, Y, df, **options) plt.clabel(cs, inline=1, fontsize=10) def savefig(filename, **options): """Save the current figure. Keyword arguments are passed along to plt.savefig https://matplotlib.org/api/_as_gen/matplotlib.pyplot.savefig.html filename: string """ print('Saving figure to file', filename) plt.savefig(filename, **options) def decorate(**options): """Decorate the current axes. Call decorate with keyword arguments like decorate(title='Title', xlabel='x', ylabel='y') The keyword arguments can be any of the axis properties https://matplotlib.org/api/axes_api.html In addition, you can use `legend=False` to suppress the legend. And you can use `loc` to indicate the location of the legend (the default value is 'best') """ loc = options.pop('loc', 'best') if options.pop('legend', True): legend(loc=loc) plt.gca().set(**options) plt.tight_layout() def legend(**options): """Draws a legend only if there is at least one labeled item. options are passed to plt.legend() https://matplotlib.org/api/_as_gen/matplotlib.pyplot.legend.html """ underride(options, loc='best') ax = plt.gca() handles, labels = ax.get_legend_handles_labels() ax.legend(handles, labels, **options) def remove_from_legend(bad_labels): """Removes some labels from the legend. bad_labels: sequence of strings """ ax = plt.gca() handles, labels = ax.get_legend_handles_labels() handle_list, label_list = [], [] for handle, label in zip(handles, labels): if label not in bad_labels: handle_list.append(handle) label_list.append(label) ax.legend(handle_list, label_list) # TODO: Either finish SubPlots or remove it class SubPlots: def __init__(self, fig, axes_seq): self.fig = fig self.axes_seq = axes_seq self.current_axes_index = 0 def current_axes(self): return self.axes_seq(self.current_axes_index) # TODO: consider making SubPlots iterable def next_axes(self): self.current_axes_index += 1 return self.current_axes() def subplots(*args, **options): fig, axes_seq = plt.subplots(*args, **options) return SubPlots(fig, axes_seq) def subplot(nrows, ncols, plot_number, **options): figsize = {(2, 1): (8, 8), (3, 1): (8, 10)} key = nrows, ncols default = (8, 5.5) width, height = figsize.get(key, default) plt.subplot(nrows, ncols, plot_number, **options) fig = plt.gcf() fig.set_figwidth(width) fig.set_figheight(height) class ModSimSeries(pd.Series): """Modified version of a Pandas Series, with a few changes to make it more suited to our purpose. In particular: 1. I provide a more consistent __init__ method. 2. Series provides two special variables called `dt` and `T` that cause problems if we try to use those names as variables. I override them so they can be used variable names. 3. Series doesn't provide a good _repr_html, so it doesn't look good in Jupyter notebooks. 4. ModSimSeries provides a set() method that takes keyword arguments. """ def __init__(self, *args, **kwargs): """Initialize a Series. Note: this cleans up a weird Series behavior, which is that Series() and Series([]) yield different results. See: https://github.com/pandas-dev/pandas/issues/16737 """ if args or kwargs: underride(kwargs, copy=True) super().__init__(*args, **kwargs) else: super().__init__([], dtype=np.float64) def _repr_html_(self): """Returns an HTML representation of the series. Mostly used for Jupyter notebooks. """ df = pd.DataFrame(self.values, index=self.index, columns=['values']) return df._repr_html_() def __copy__(self, deep=True): series = super().copy(deep=deep) return self.__class__(series) copy = __copy__ def set(self, **kwargs): """Uses keyword arguments to update the Series in place. Example: series.set(a=1, b=2) """ for name, value in kwargs.items(): self[name] = value @property def dt(self): """Intercept the Series accessor object so we can use `dt` as a row label and access it using dot notation. https://pandas.pydata.org/pandas-docs/stable/generated/pandas.Series.dt.html """ return self.loc['dt'] @property def T(self): """Intercept the Series accessor object so we can use `T` as a row label and access it using dot notation. https://pandas.pydata.org/pandas-docs/stable/generated/pandas.Series.T.html """ return self.loc['T'] def get_first_label(series): """Returns the label of the first element.""" return series.index[0] def get_last_label(series): """Returns the label of the first element.""" return series.index[-1] def get_index_label(series, i): """Returns the ith label in the index.""" return series.index[i] def get_first_value(series): """Returns the value of the first element.""" return series.values[0] def get_last_value(series): """Returns the value of the first element.""" return series.values[-1] def gradient(series): """Computes the numerical derivative of a series.""" a = np.gradient(series, series.index) return TimeSeries(a, series.index) class TimeSeries(ModSimSeries): """Represents a mapping from times to values.""" pass class SweepSeries(ModSimSeries): """Represents a mapping from parameter values to metrics.""" pass class System(ModSimSeries): """Contains system variables and their values. Takes keyword arguments and stores them as rows. """ def __init__(self, *args, **kwargs): """Initialize the series. If there are no positional arguments, use kwargs. If there is one positional argument, copy it and add in the kwargs. More than one positional argument is an error. """ if len(args) == 0: super().__init__(list(kwargs.values()), index=kwargs) elif len(args) == 1: super().__init__(*args, copy=True) self.set(**kwargs) else: msg = '__init__() takes at most one positional argument' raise TypeError(msg) class State(System): """Contains state variables and their values. Takes keyword arguments and stores them as rows. """ pass class Condition(System): """Represents the condition of a system. Condition objects are often used to construct a System object. """ pass class Params(System): """Represents a set of parameters. """ pass def compute_abs_diff(seq): xs = np.asarray(seq) diff = np.ediff1d(xs, np.nan) if isinstance(seq, Series): return Series(diff, seq.index) else: return diff def compute_rel_diff(seq): xs = np.asarray(seq) diff = np.ediff1d(xs, np.nan) return diff / seq class ModSimDataFrame(pd.DataFrame): """ModSimDataFrame is a modified version of a Pandas DataFrame, with a few changes to make it more suited to our purpose. In particular: 1. DataFrame provides two special variables called `dt` and `T` that cause problems if we try to use those names as variables. I override them so they can be used as row labels. 2. When you select a row or column from a ModSimDataFrame, you get back an appropriate subclass of Series: TimeSeries, SweepSeries, or ModSimSeries. """ column_constructor = ModSimSeries row_constructor = ModSimSeries def __init__(self, *args, **options): # TODO: currently ModSimDataFrame underrides to float64 and # ModSimSeries does not. Does this inconsistency make sense? # underride(options, dtype=np.float64) super().__init__(*args, **options) def __getitem__(self, key): """Intercept the column getter to return the right subclass of Series. """ obj = super().__getitem__(key) if isinstance(obj, Series): obj = self.column_constructor(obj) return obj @property def dt(self): """Intercept the Series accessor object so we can use `dt` as a column label and access it using dot notation. https://pandas.pydata.org/pandas-docs/stable/generated/pandas.DataFrame.dt.html """ return self['dt'] @property def T(self): """Intercept the Series accessor object so we can use `T` as a column label and access it using dot notation. https://pandas.pydata.org/pandas-docs/stable/generated/pandas.DataFrame.T.html """ return self['T'] @property def row(self): """Gets or sets a row. Returns a wrapper for the Pandas LocIndexer, so when we look up a row we get the right kind of ModSimSeries. returns ModSimLocIndexer """ li = self.loc return ModSimLocIndexer(li, self.row_constructor) class ModSimLocIndexer: """Wraps a Pandas LocIndexer.""" def __init__(self, li, constructor): """Save the LocIndexer and constructor. """ self.li = li self.constructor = constructor def __getitem__(self, key): """Get a row and return the appropriate type of Series. """ result = self.li[key] if isinstance(result, Series): result = self.constructor(result) return result def __setitem__(self, key, value): """Setting just passes the request to the wrapped object. """ self.li[key] = value class TimeFrame(ModSimDataFrame): """A DataFrame that maps from time to State. """ column_constructor = TimeSeries row_constructor = State class SweepFrame(ModSimDataFrame): """A DataFrame that maps from a parameter value to a SweepSeries. """ column_constructor = SweepSeries row_constructor = SweepSeries def Vector(*args, units=None): """Make a ModSimVector. args: can be a single argument or sequence units: Pint Unit object or Quantity If there's only one argument, it should be a sequence. Otherwise, the arguments are treated as coordinates. returns: ModSimVector """ if len(args) == 1: args = args[0] # if it's a series, pull out the values if isinstance(args, Series): args = args.values # see if any of the arguments have units; if so, save the first one for elt in args: found_units = getattr(elt, 'units', None) if found_units: break if found_units: # if there are units, remove them args = [magnitude(elt) for elt in args] # if the units keyword is provided, it overrides the units in args if units is not None: found_units = units return ModSimVector(args, found_units) ## Vector functions (should work with any sequence) def vector_mag(v): """Vector magnitude with units. returns: number or Quantity """ return np.sqrt(np.dot(v, v)) * units(v) def vector_mag2(v): """Vector magnitude squared with units. returns: number of Quantity """ return np.dot(v, v) * units(v) * units(v) def vector_angle(v): """Angle between v and the positive x axis. Only works with 2-D vectors. returns: number in radians """ assert len(v) == 2 x, y = v return np.arctan2(y, x) def vector_polar(v): """Vector magnitude and angle. returns: (number or quantity, number in radians) """ return vector_mag(v), vector_angle(v) def vector_hat(v): """Unit vector in the direction of v. The result should have no units. returns: Vector or array """ # get the size of the vector mag = vector_mag(v) # check if the magnitude of the Quantity is 0 if magnitude(mag) == 0: if isinstance(v, ModSimVector): return Vector(magnitude(v)) else: return magnitude(np.asarray(v)) else: return v / mag def vector_perp(v): """Perpendicular Vector (rotated left). Only works with 2-D Vectors. returns: Vector """ assert len(v) == 2 x, y = v return Vector(-y, x) def vector_dot(v, w): """Dot product of v and w. returns: number or Quantity """ return np.dot(v, w) * units(v) * units(w) def vector_cross(v, w): """Cross product of v and w. returns: number or Quantity for 2-D, Vector for 3-D """ res = np.cross(v, w) if len(v)==3 and (isinstance(v, ModSimVector) or isinstance(w, ModSimVector)): return ModSimVector(res, units(v) * units(w)) else: return res * units(v) * units(w) def vector_proj(v, w): """Projection of v onto w. Results has the units of v, but that might not make sense unless v and w have the same units. returns: array or Vector with direction of w and units of v. """ w_hat = vector_hat(w) return vector_dot(v, w_hat) * w_hat def scalar_proj(v, w): """Returns the scalar projection of v onto w. Which is the magnitude of the projection of v onto w. Results has the units of v, but that might not make sense unless v and w have the same units. returns: scalar with units of v. """ return vector_dot(v, vector_hat(w)) def vector_dist(v, w): """Euclidean distance from v to w, with units.""" if isinstance(v, list): v = np.asarray(v) return vector_mag(v-w) def vector_diff_angle(v, w): """Angular difference between two vectors, in radians. """ if len(v) == 2: return vector_angle(v) - vector_angle(w) else: #TODO: see http://www.euclideanspace.com/maths/algebra/ # vectors/angleBetween/ raise NotImplementedError() class ModSimVector(Quantity): """Represented as a Pint Quantity with a NumPy array x, y, z, mag, mag2, and angle are accessible as attributes. """ @property def x(self): """Returns the x component with units.""" return self[0] @property def y(self): """Returns the y component with units.""" return self[1] @property def z(self): """Returns the z component with units.""" return self[2] @property def mag(self): """Returns the magnitude with units.""" return vector_mag(self) @property def mag2(self): """Returns the magnitude squared with units.""" return vector_mag2(self) @property def angle(self): """Returns the angle between self and the positive x axis.""" return vector_angle(self) # make the vector functions available as methods polar = vector_polar hat = vector_hat perp = vector_perp dot = vector_dot cross = vector_cross proj = vector_proj comp = scalar_proj dist = vector_dist diff_angle = vector_diff_angle def plot_segment(A, B, **options): """Plots a line segment between two Vectors. For 3-D vectors, the z axis is ignored. Additional options are passed along to plot(). A: Vector B: Vector """ xs = A.x, B.x ys = A.y, B.y plot(xs, ys, **options) @property def dimensionality(self): """Unit's dimensionality (e.g. {length: 1, time: -1}) This is a simplified version of this method that does no caching. returns: dimensionality """ dim = self._REGISTRY._get_dimensionality(self._units) return dim # monkey patch Unit and Quantity so they use the non-caching # version of `dimensionality` pint.unit._Unit.dimensionality = dimensionality pint.quantity._Quantity.dimensionality = dimensionality class units_off: SAVED_PINT_METHOD_STACK = [] def __enter__(self): """Make all quantities behave as if they were dimensionless. """ self.SAVED_PINT_METHOD_STACK.append(UNITS._get_dimensionality) UNITS._get_dimensionality = lambda self: {} def __exit__(self, type, value, traceback): """Restore the saved behavior of quantities. """ UNITS._get_dimensionality = self.SAVED_PINT_METHOD_STACK.pop()