# Source code for vaex.geo

import vaex
import numpy as np
from .utils import _ensure_strings_from_expressions, _ensure_string_from_expression, as_flat_array
from .dataframe import docsubst

[docs]class DataFrameAccessorGeo(object): """Geometry/geographic helper methods Example: >>> df_xyz = df.geo.spherical2cartesian(df.longitude, df.latitude, df.distance) >>> df_xyz.x.mean() """
[docs] def __init__(self, df): self.df = df
[docs] @docsubst def spherical2cartesian(self, alpha, delta, distance, xname="x", yname="y", zname="z", propagate_uncertainties=False, center=[0, 0, 0], radians=False, inplace=False): """Convert spherical to cartesian coordinates. :param alpha: :param delta: polar angle, ranging from the -90 (south pole) to 90 (north pole) :param distance: radial distance, determines the units of x, y and z :param xname: :param yname: :param zname: :param propagate_uncertainties: {propagate_uncertainties} :param center: :param radians: :return: New dataframe (in inplace is False) with new x,y,z columns """ df = self.df if inplace else self.df.copy() alpha = df._expr(alpha) delta = df._expr(delta) distance = df._expr(distance) if not radians: df.add_variable('pi', np.pi) alpha = alpha * df._expr('pi')/180 delta = delta * df._expr('pi')/180 # TODO: use sth like .optimize by default to get rid of the +0 ? if center[0]: df[xname] = np.cos(alpha) * np.cos(delta) * distance + center[0] else: df[xname] = np.cos(alpha) * np.cos(delta) * distance if center[1]: df[yname] = np.sin(alpha) * np.cos(delta) * distance + center[1] else: df[yname] = np.sin(alpha) * np.cos(delta) * distance if center[2]: df[zname] = np.sin(delta) * distance + center[2] else: df[zname] = np.sin(delta) * distance if propagate_uncertainties: df.propagate_uncertainties([df[xname], df[yname], df[zname]]) return df
[docs] def cartesian2spherical(self, x="x", y="y", z="z", alpha="l", delta="b", distance="distance", radians=False, center=None, center_name="solar_position", inplace=False): """Convert cartesian to spherical coordinates. :param x: :param y: :param z: :param alpha: :param delta: name for polar angle, ranges from -90 to 90 (or -pi to pi when radians is True). :param distance: :param radians: :param center: :param center_name: :return: """ df = self.df if inplace else self.df.copy() df.add_variable('pi', np.pi) transform = "" if radians else "*180./pi" if center is not None: df.add_variable(center_name, center) if center is not None and center[0] != 0: x = "({x} - {center_name}[0])".format(**locals()) if center is not None and center[1] != 0: y = "({y} - {center_name}[1])".format(**locals()) if center is not None and center[2] != 0: z = "({z} - {center_name}[2])".format(**locals()) df.add_virtual_column(distance, "sqrt({x}**2 + {y}**2 + {z}**2)".format(**locals())) # df.add_virtual_column(alpha, "((arctan2({y}, {x}) + 2*pi) % (2*pi)){transform}".format(**locals())) df.add_virtual_column(alpha, "arctan2({y}, {x}){transform}".format(**locals())) df.add_virtual_column(delta, "(-arccos({z}/{distance})+pi/2){transform}".format(**locals())) return df