Coverage for python/lsst/dax/apdb/tests/data_factory.py: 23%
49 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-11 03:30 -0700
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-11 03:30 -0700
1# This file is part of dax_apdb.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <http://www.gnu.org/licenses/>.
22from __future__ import annotations
24import random
25from collections.abc import Iterator
26from typing import Any
28import astropy.time
29import numpy
30import pandas
31from lsst.sphgeom import LonLat, Region, UnitVector3d
34def _genPointsInRegion(region: Region, count: int) -> Iterator[LonLat]:
35 """Generate bunch of SpherePoints inside given region.
37 Parameters
38 ----------
39 region : `lsst.sphgeom.Region`
40 Spherical region.
41 count : `int`
42 Number of points to generate.
44 Notes
45 -----
46 Returned points are random but not necessarily uniformly distributed.
47 """
48 bbox = region.getBoundingBox()
49 center = bbox.getCenter()
50 center_lon = center.getLon().asRadians()
51 center_lat = center.getLat().asRadians()
52 width = bbox.getWidth().asRadians()
53 height = bbox.getHeight().asRadians()
54 while count > 0:
55 lon = random.uniform(center_lon - width / 2, center_lon + width / 2)
56 lat = random.uniform(center_lat - height / 2, center_lat + height / 2)
57 lonlat = LonLat.fromRadians(lon, lat)
58 uv3d = UnitVector3d(lonlat)
59 if region.contains(uv3d):
60 yield lonlat
61 count -= 1
64def makeObjectCatalog(
65 region: Region, count: int, visit_time: astropy.time.Time, *, start_id: int = 1, **kwargs: Any
66) -> pandas.DataFrame:
67 """Make a catalog containing a bunch of DiaObjects inside a region.
69 Parameters
70 ----------
71 region : `lsst.sphgeom.Region`
72 Spherical region.
73 count : `int`
74 Number of records to generate.
75 visit_time : `astropy.time.Time`
76 Time of the visit.
77 start_id : `int`
78 Starting diaObjectId.
79 **kwargs : `Any`
80 Additional columns and their values to add to catalog.
82 Returns
83 -------
84 catalog : `pandas.DataFrame`
85 Catalog of DiaObjects records.
87 Notes
88 -----
89 Returned catalog only contains three columns - ``diaObjectId`, ``ra``, and
90 ``dec`` (in degrees).
91 """
92 points = list(_genPointsInRegion(region, count))
93 # diaObjectId=0 may be used in some code for DiaSource foreign key to mean
94 # the same as ``None``.
95 ids = numpy.arange(start_id, len(points) + start_id, dtype=numpy.int64)
96 ras = numpy.array([lonlat.getLon().asDegrees() for lonlat in points], dtype=numpy.float64)
97 decs = numpy.array([lonlat.getLat().asDegrees() for lonlat in points], dtype=numpy.float64)
98 nDiaSources = numpy.ones(len(points), dtype=numpy.int32)
99 dt = visit_time.datetime
100 data = dict(
101 kwargs,
102 diaObjectId=ids,
103 ra=ras,
104 dec=decs,
105 nDiaSources=nDiaSources,
106 lastNonForcedSource=dt,
107 )
108 df = pandas.DataFrame(data)
109 return df
112def makeSourceCatalog(
113 objects: pandas.DataFrame,
114 visit_time: astropy.time.Time,
115 start_id: int = 0,
116 visit: int = 1,
117 detector: int = 1,
118) -> pandas.DataFrame:
119 """Make a catalog containing a bunch of DiaSources associated with the
120 input DiaObjects.
122 Parameters
123 ----------
124 objects : `pandas.DataFrame`
125 Catalog of DiaObject records.
126 visit_time : `astropy.time.Time`
127 Time of the visit.
128 start_id : `int`
129 Starting value for ``diaObjectId``.
130 visit, detector : `int`
131 Value for ``visit`` and ``detector`` fields.
133 Returns
134 -------
135 catalog : `pandas.DataFrame`
136 Catalog of DiaSource records.
138 Notes
139 -----
140 Returned catalog only contains small number of columns needed for tests.
141 """
142 nrows = len(objects)
143 midpointMjdTai = visit_time.mjd
144 df = pandas.DataFrame(
145 {
146 "diaSourceId": numpy.arange(start_id, start_id + nrows, dtype=numpy.int64),
147 "diaObjectId": objects["diaObjectId"],
148 "visit": numpy.full(nrows, visit, dtype=numpy.int64),
149 "detector": numpy.full(nrows, detector, dtype=numpy.int16),
150 "parentDiaSourceId": 0,
151 "ra": objects["ra"],
152 "dec": objects["dec"],
153 "midpointMjdTai": numpy.full(nrows, midpointMjdTai, dtype=numpy.float64),
154 "flags": numpy.full(nrows, 0, dtype=numpy.int64),
155 }
156 )
157 return df
160def makeForcedSourceCatalog(
161 objects: pandas.DataFrame, visit_time: astropy.time.Time, visit: int = 1, detector: int = 1
162) -> pandas.DataFrame:
163 """Make a catalog containing a bunch of DiaForcedSources associated with
164 the input DiaObjects.
166 Parameters
167 ----------
168 objects : `pandas.DataFrame`
169 Catalog of DiaObject records.
170 visit_time : `astropy.time.Time`
171 Time of the visit.
172 visit, detector : `int`
173 Value for ``visit`` and ``detector`` fields.
175 Returns
176 -------
177 catalog : `pandas.DataFrame`
178 Catalog of DiaForcedSource records.
180 Notes
181 -----
182 Returned catalog only contains small number of columns needed for tests.
183 """
184 nrows = len(objects)
185 midpointMjdTai = visit_time.mjd
186 df = pandas.DataFrame(
187 {
188 "diaObjectId": objects["diaObjectId"],
189 "visit": numpy.full(nrows, visit, dtype=numpy.int64),
190 "detector": numpy.full(nrows, detector, dtype=numpy.int16),
191 "midpointMjdTai": numpy.full(nrows, midpointMjdTai, dtype=numpy.float64),
192 "flags": numpy.full(nrows, 0, dtype=numpy.int64),
193 }
194 )
195 return df
198def makeSSObjectCatalog(count: int, start_id: int = 1, flags: int = 0) -> pandas.DataFrame:
199 """Make a catalog containing a bunch of SSObjects.
201 Parameters
202 ----------
203 count : `int`
204 Number of records to generate.
205 startID : `int`
206 Initial SSObject ID.
207 flags : `int`
208 Value for ``flags`` column.
210 Returns
211 -------
212 catalog : `pandas.DataFrame`
213 Catalog of SSObjects records.
215 Notes
216 -----
217 Returned catalog only contains three columns - ``ssObjectId`, ``arc``,
218 and ``flags``.
219 """
220 ids = numpy.arange(start_id, count + start_id, dtype=numpy.int64)
221 arc = numpy.full(count, 0.001, dtype=numpy.float32)
222 flags_array = numpy.full(count, flags, dtype=numpy.int64)
223 df = pandas.DataFrame({"ssObjectId": ids, "arc": arc, "flags": flags_array})
224 return df