6733525095

Committed 02 Nov 2023 01:52PM UTC coverage: 40.544% (+0.9%) from 39.595%

Build # 6733525095

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #123 from moeyensj/v2.0-link-aims-sample

Link AIMS sample

Run Details

301 of 301 new or added lines in 12 files covered. (100.0%)

1878 of 4632 relevant lines covered (40.54%)

0.41 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

75.71

/thor/projections/gnomonic.py

from typing import Literal, Optional

import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.compute as pc
import quivr as qv
from adam_core.coordinates import CartesianCoordinates, Origin
from adam_core.coordinates.covariances import transform_covariances_jacobian
from adam_core.time import Timestamp
from typing_extensions import Self

from .covariances import ProjectionCovariances
from .transforms import _cartesian_to_gnomonic, cartesian_to_gnomonic


class GnomonicCoordinates(qv.Table):

    time = Timestamp.as_column(nullable=True)
    theta_x = qv.Float64Column()
    theta_y = qv.Float64Column()
    vtheta_x = qv.Float64Column(nullable=True)
    vtheta_y = qv.Float64Column(nullable=True)
    covariance = ProjectionCovariances.as_column(nullable=True)
    origin = Origin.as_column()
    frame = qv.StringAttribute("testorbit?")

    @property
    def values(self) -> np.ndarray:
        return np.array(
            self.table.select(["theta_x", "theta_y", "vtheta_x", "vtheta_y"])
        )

    @property
    def sigma_theta_x(self) -> np.ndarray:
        """
        1-sigma uncertainty in the theta_x coordinate.
        """
        return self.covariance.sigmas[:, 0]

    @property
    def sigma_theta_y(self) -> np.ndarray:
        """
        1-sigma uncertainty in the theta_y coordinate.
        """
        return self.covariance.sigmas[:, 1]

    @property
    def sigma_vtheta_x(self) -> np.ndarray:
        """
        1-sigma uncertainty in the theta_x coordinate velocity.
        """
        return self.covariance.sigmas[:, 2]

    @property
    def sigma_vtheta_y(self):
        """
        1-sigma uncertainty in the theta_y coordinate velocity.
        """
        return self.covariance.sigmas[:, 3]

    @classmethod
    def from_cartesian(
        cls,
        cartesian: CartesianCoordinates,
        center_cartesian: Optional[CartesianCoordinates] = None,
    ) -> Self:
        """
        Create a GnomonicCoordinates object from a CartesianCoordinates object.

        Parameters
        ----------
        cartesian : `~adam_core.coordinates.cartesian.CartesianCoordinates`
            Cartesian coordinates to be transformed.
        center_cartesian : `CartesianCoordinates`, optional
            Cartesian coordinates of the center of the projection. Default is the
            x-axis at 1 au.

        Returns
        -------
        gnomonic : `~thor.projections.gnomonic.GnomonicCoordinates`
            Gnomonic coordinates.
        """
        assert len(cartesian.origin.code.unique()) == 1

        # Check if input coordinates have times defined, if they do lets make
        # sure that we have the correct cartesian center coordinates for each
        if center_cartesian is None and cartesian.time is None:
            # Both center and input coordinates have no times defined, so we
            # can just use the default center coordinates
            center_cartesian = CartesianCoordinates.from_kwargs(
                x=[1],
                y=[0],
                z=[0],
                vx=[0],
                vy=[0],
                vz=[0],
                frame=cartesian.frame,
                origin=cartesian[0].origin,
            )

            # Set the linkage key to be the same integers for each cartesian
            # coordinate and center cartesian coordinate
            left_key = pa.array(np.zeros(len(cartesian)), type=pa.int64())
            right_key = pa.array(np.zeros(len(center_cartesian)), type=pa.int64())

            # Create a linkage between the cartesian coordinates and the center cartesian
            # coordinates on time
            link = qv.Linkage(
                cartesian,
                center_cartesian,
                left_keys=left_key,
                right_keys=right_key,
            )

        elif center_cartesian is None and cartesian.time is not None:
            # Create a center cartesian coordinate for each unique time in the
            # input cartesian coordinates
            times = cartesian.time.jd().unique()
            num_unique_times = len(times)
            center_cartesian = CartesianCoordinates.from_kwargs(
                time=cartesian.time.unique(),
                x=np.ones(num_unique_times, dtype=np.float64),
                y=np.zeros(num_unique_times, dtype=np.float64),
                z=np.zeros(num_unique_times, dtype=np.float64),
                vx=np.zeros(num_unique_times, dtype=np.float64),
                vy=np.zeros(num_unique_times, dtype=np.float64),
                vz=np.zeros(num_unique_times, dtype=np.float64),
                frame=cartesian.frame,
                origin=qv.concatenate(
                    [cartesian[0].origin for _ in range(num_unique_times)]
                ),
            )
            link = cartesian.time.link(center_cartesian.time, precision="ms")

        elif center_cartesian is not None and cartesian.time is not None:
            assert cartesian.time.scale == center_cartesian.time.scale

            link = cartesian.time.link(center_cartesian.time, precision="ms")

        else:
            raise ValueError(
                "Input cartesian coordinates and projection center cartesian coordinates "
                "must have the same times."
            )

        # Round the times to the nearest microsecond and use those to select
        # the cartesian coordinates and center cartesian coordinates
        rounded_cartesian_times = cartesian.time.rounded(precision="ms")  # type: ignore
        rounded_center_cartesian_times = center_cartesian.time.rounded(precision="ms")  # type: ignore

        gnomonic_coords = []
        for key, time_i, center_time_i in link.iterate():

            cartesian_i = cartesian.apply_mask(
                rounded_cartesian_times.equals_scalar(key[0], key[1], precision="ms")
            )
            center_cartesian_i = center_cartesian.apply_mask(
                rounded_center_cartesian_times.equals_scalar(
                    key[0], key[1], precision="ms"
                )
            )

            if len(center_cartesian_i) == 0:
                raise ValueError("No center cartesian coordinates found for this time.")

            coords_gnomonic, M = cartesian_to_gnomonic(
                cartesian_i.values,
                center_cartesian=center_cartesian_i.values[0],
            )
            coords_gnomonic = np.array(coords_gnomonic)

            if not cartesian_i.covariance.is_all_nan():
                cartesian_covariances = cartesian_i.covariance.to_matrix()
                covariances_gnomonic = transform_covariances_jacobian(
                    cartesian_i.values,
                    cartesian_covariances,
                    _cartesian_to_gnomonic,
                    in_axes=(0, None),
                    center_cartesian=center_cartesian_i.values[0],
                )
            else:
                covariances_gnomonic = np.empty(
                    (len(coords_gnomonic), 4, 4), dtype=np.float64
                )
                covariances_gnomonic.fill(np.nan)

            gnomonic_coords.append(
                cls.from_kwargs(
                    theta_x=coords_gnomonic[:, 0],
                    theta_y=coords_gnomonic[:, 1],
                    vtheta_x=coords_gnomonic[:, 2],
                    vtheta_y=coords_gnomonic[:, 3],
                    time=cartesian_i.time,
                    covariance=ProjectionCovariances.from_matrix(covariances_gnomonic),
                    origin=cartesian_i.origin,
                    frame=cartesian_i.frame,
                )
            )

        return qv.concatenate(gnomonic_coords)

1	from typing import Literal, Optional	1✔
2
3	import numpy as np	1✔
4	import pandas as pd	1✔
5	import pyarrow as pa	1✔
6	import pyarrow.compute as pc	1✔
7	import quivr as qv	1✔
8	from adam_core.coordinates import CartesianCoordinates, Origin	1✔
9	from adam_core.coordinates.covariances import transform_covariances_jacobian	1✔
10	from adam_core.time import Timestamp	1✔
11	from typing_extensions import Self	1✔
12
13	from .covariances import ProjectionCovariances	1✔
14	from .transforms import _cartesian_to_gnomonic, cartesian_to_gnomonic	1✔
15
16
17	class GnomonicCoordinates(qv.Table):	1✔
18
19	time = Timestamp.as_column(nullable=True)	1✔
20	theta_x = qv.Float64Column()	1✔
21	theta_y = qv.Float64Column()	1✔
22	vtheta_x = qv.Float64Column(nullable=True)	1✔
23	vtheta_y = qv.Float64Column(nullable=True)	1✔
24	covariance = ProjectionCovariances.as_column(nullable=True)	1✔
25	origin = Origin.as_column()	1✔
26	frame = qv.StringAttribute("testorbit?")	1✔
27
28	@property	1✔
29	def values(self) -> np.ndarray:	1✔
30	return np.array(	×
31	self.table.select(["theta_x", "theta_y", "vtheta_x", "vtheta_y"])
32	)
33
34	@property	1✔
35	def sigma_theta_x(self) -> np.ndarray:	1✔
36	"""
37	1-sigma uncertainty in the theta_x coordinate.
38	"""
39	return self.covariance.sigmas[:, 0]	×
40
41	@property	1✔
42	def sigma_theta_y(self) -> np.ndarray:	1✔
43	"""
44	1-sigma uncertainty in the theta_y coordinate.
45	"""
46	return self.covariance.sigmas[:, 1]	×
47
48	@property	1✔
49	def sigma_vtheta_x(self) -> np.ndarray:	1✔
50	"""
51	1-sigma uncertainty in the theta_x coordinate velocity.
52	"""
53	return self.covariance.sigmas[:, 2]	×
54
55	@property	1✔
56	def sigma_vtheta_y(self):	1✔
57	"""
58	1-sigma uncertainty in the theta_y coordinate velocity.
59	"""
60	return self.covariance.sigmas[:, 3]	×
61
62	@classmethod	1✔
63	def from_cartesian(	1✔
64	cls,
65	cartesian: CartesianCoordinates,
66	center_cartesian: Optional[CartesianCoordinates] = None,
67	) -> Self:
68	"""
69	Create a GnomonicCoordinates object from a CartesianCoordinates object.
70
71	Parameters
72	----------
73	cartesian : `~adam_core.coordinates.cartesian.CartesianCoordinates`
74	Cartesian coordinates to be transformed.
75	center_cartesian : `CartesianCoordinates`, optional
76	Cartesian coordinates of the center of the projection. Default is the
77	x-axis at 1 au.
78
79	Returns
80	-------
81	gnomonic : `~thor.projections.gnomonic.GnomonicCoordinates`
82	Gnomonic coordinates.
83	"""
84	assert len(cartesian.origin.code.unique()) == 1	1✔
85
86	# Check if input coordinates have times defined, if they do lets make
87	# sure that we have the correct cartesian center coordinates for each
88	if center_cartesian is None and cartesian.time is None:	1✔
89	# Both center and input coordinates have no times defined, so we
90	# can just use the default center coordinates
91	center_cartesian = CartesianCoordinates.from_kwargs(	×
92	x=[1],
93	y=[0],
94	z=[0],
95	vx=[0],
96	vy=[0],
97	vz=[0],
98	frame=cartesian.frame,
99	origin=cartesian[0].origin,
100	)
101
102	# Set the linkage key to be the same integers for each cartesian
103	# coordinate and center cartesian coordinate
104	left_key = pa.array(np.zeros(len(cartesian)), type=pa.int64())	×
105	right_key = pa.array(np.zeros(len(center_cartesian)), type=pa.int64())	×
106
107	# Create a linkage between the cartesian coordinates and the center cartesian
108	# coordinates on time
109	link = qv.Linkage(	×
110	cartesian,
111	center_cartesian,
112	left_keys=left_key,
113	right_keys=right_key,
114	)
115
116	elif center_cartesian is None and cartesian.time is not None:	1✔
117	# Create a center cartesian coordinate for each unique time in the
118	# input cartesian coordinates
119	times = cartesian.time.jd().unique()	×
120	num_unique_times = len(times)	×
121	center_cartesian = CartesianCoordinates.from_kwargs(	×
122	time=cartesian.time.unique(),
123	x=np.ones(num_unique_times, dtype=np.float64),
124	y=np.zeros(num_unique_times, dtype=np.float64),
125	z=np.zeros(num_unique_times, dtype=np.float64),
126	vx=np.zeros(num_unique_times, dtype=np.float64),
127	vy=np.zeros(num_unique_times, dtype=np.float64),
128	vz=np.zeros(num_unique_times, dtype=np.float64),
129	frame=cartesian.frame,
130	origin=qv.concatenate(
131	[cartesian[0].origin for _ in range(num_unique_times)]
132	),
133	)
134	link = cartesian.time.link(center_cartesian.time, precision="ms")	×
135
136	elif center_cartesian is not None and cartesian.time is not None:	1✔
137	assert cartesian.time.scale == center_cartesian.time.scale	1✔
138
139	link = cartesian.time.link(center_cartesian.time, precision="ms")	1✔
140
141	else:
142	raise ValueError(	×
143	"Input cartesian coordinates and projection center cartesian coordinates "
144	"must have the same times."
145	)
146
147	# Round the times to the nearest microsecond and use those to select
148	# the cartesian coordinates and center cartesian coordinates
149	rounded_cartesian_times = cartesian.time.rounded(precision="ms") # type: ignore	1✔
150	rounded_center_cartesian_times = center_cartesian.time.rounded(precision="ms") # type: ignore	1✔
151
152	gnomonic_coords = []	1✔
153	for key, time_i, center_time_i in link.iterate():	1✔
154
155	cartesian_i = cartesian.apply_mask(	1✔
156	rounded_cartesian_times.equals_scalar(key[0], key[1], precision="ms")
157	)
158	center_cartesian_i = center_cartesian.apply_mask(	1✔
159	rounded_center_cartesian_times.equals_scalar(
160	key[0], key[1], precision="ms"
161	)
162	)
163
164	if len(center_cartesian_i) == 0:	1✔
165	raise ValueError("No center cartesian coordinates found for this time.")	×
166
167	coords_gnomonic, M = cartesian_to_gnomonic(	1✔
168	cartesian_i.values,
169	center_cartesian=center_cartesian_i.values[0],
170	)
171	coords_gnomonic = np.array(coords_gnomonic)	1✔
172
173	if not cartesian_i.covariance.is_all_nan():	1✔
174	cartesian_covariances = cartesian_i.covariance.to_matrix()	×
175	covariances_gnomonic = transform_covariances_jacobian(	×
176	cartesian_i.values,
177	cartesian_covariances,
178	_cartesian_to_gnomonic,
179	in_axes=(0, None),
180	center_cartesian=center_cartesian_i.values[0],
181	)
182	else:
183	covariances_gnomonic = np.empty(	1✔
184	(len(coords_gnomonic), 4, 4), dtype=np.float64
185	)
186	covariances_gnomonic.fill(np.nan)	1✔
187
188	gnomonic_coords.append(	1✔
189	cls.from_kwargs(
190	theta_x=coords_gnomonic[:, 0],
191	theta_y=coords_gnomonic[:, 1],
192	vtheta_x=coords_gnomonic[:, 2],
193	vtheta_y=coords_gnomonic[:, 3],
194	time=cartesian_i.time,
195	covariance=ProjectionCovariances.from_matrix(covariances_gnomonic),
196	origin=cartesian_i.origin,
197	frame=cartesian_i.frame,
198	)
199	)
200
201	return qv.concatenate(gnomonic_coords)	1✔

moeyensj / thor / 6733525095

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous