Skip to content

Commit

Permalink
Improve bookkeeping by adding name to SmallBodies and Pointings and u…
Browse files Browse the repository at this point in the history
…sing it to create file names for output data
  • Loading branch information
moeyensj committed Sep 19, 2024
1 parent c3a536a commit ddc2518
Show file tree
Hide file tree
Showing 6 changed files with 26 additions and 9 deletions.
3 changes: 2 additions & 1 deletion src/adam_test_data/datasets/s3m.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@ def load_S3M(directory: str) -> SmallBodies:
"""
s3m_files = sorted(glob.glob(os.path.join(directory, "*.s3m")))

s3m = SmallBodies.empty()
s3m = SmallBodies.empty(name="S3M")
id_offset = 0
for s3m_file in s3m_files:
print(f"Loading {s3m_file}")
Expand Down Expand Up @@ -115,6 +115,7 @@ def load_S3M(directory: str) -> SmallBodies:
s3m_i = SmallBodies.from_kwargs(
orbits=orbits_i,
properties=photometric_properties_i,
name="S3M",
)

s3m = qv.concatenate([s3m, s3m_i])
Expand Down
16 changes: 11 additions & 5 deletions src/adam_test_data/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -595,7 +595,7 @@ def sorcha_worker(
)

# Create a subdirectory for this chunk
chunk_base = f"chunk_{orbit_ids_indices[0]:08d}_{orbit_ids_indices[1]:08d}"
chunk_base = f"{orbit_ids_indices[0]:08d}_{orbit_ids_indices[1]:08d}"
output_dir_chunk = os.path.join(output_dir, chunk_base)

catalog = sorcha(
Expand All @@ -612,7 +612,7 @@ def sorcha_worker(
)

# Serialize the output tables to parquet and return the paths
catalog_file = os.path.join(output_dir, f"{chunk_base}_{tag}.parquet")
catalog_file = os.path.join(output_dir, f"{tag}_{chunk_base}.parquet")
catalog.to_parquet(catalog_file)

return catalog_file
Expand Down Expand Up @@ -773,7 +773,7 @@ def generate_test_data(
observatory: Observatory,
noise_densities: Optional[list[float]] = None,
time_range: Optional[list[float]] = None,
tag: str = "sorcha",
tag: Optional[str] = None,
overwrite: bool = True,
randomization: bool = True,
output_columns: Literal["basic", "all"] = "all",
Expand Down Expand Up @@ -802,7 +802,8 @@ def generate_test_data(
time_range : list[float], optional
The time range to filter the pointings by, by default None.
tag : str, optional
The tag to use for the output files, by default "sorcha".
The tag to use for the output files and the catalog ID. If None, the tag will be
generated from the pointings name, observatory code, and small bodies name.
randomization : bool, optional
Ramdomize the photometry and astrometry using the calculated uncertainties.
output_columns : Literal["basic", "all"], optional
Expand Down Expand Up @@ -837,6 +838,9 @@ def generate_test_data(
# Create the output directory if it doesn't exist
os.makedirs(output_dir, exist_ok=True)

if tag is None:
tag = f"{pointings.name}_{observatory.code}_{small_bodies.name}"

# Run sorcha
catalog_file = run_sorcha(
output_dir,
Expand All @@ -862,13 +866,15 @@ def generate_test_data(

# Now generate noise observations
for noise_density in noise_densities:
tag_noise = f"{tag}_noise_{noise_density:.3f}"

# Add noise to the observations
noise_catalog = generate_noise(
output_dir,
pointings_filtered,
observatory,
noise_density,
tag=tag,
tag=tag_noise,
seed=seed,
chunk_size=chunk_size,
max_processes=max_processes,
Expand Down
10 changes: 8 additions & 2 deletions src/adam_test_data/noise.py
Original file line number Diff line number Diff line change
Expand Up @@ -326,9 +326,12 @@ def noise_worker(
seed=seed,
)

if tag is None:
tag = f"noise_{density:.3f}"

noise_file = os.path.join(
out_dir,
f"noise_{pointing_ids_indices[0]:08d}_{pointing_ids_indices[1]:08d}.parquet",
f"{tag}_{pointing_ids_indices[0]:08d}_{pointing_ids_indices[1]:08d}.parquet",
)
noise.to_parquet(noise_file)
return noise_file
Expand Down Expand Up @@ -386,7 +389,10 @@ def generate_noise(
# Create the output directory if it doesn't exist
os.makedirs(output_dir, exist_ok=True)

noise_file = os.path.join(output_dir, f"noise_{density:.3f}.parquet")
if tag is None:
tag = f"noise_{density:.3f}"

noise_file = os.path.join(output_dir, f"{tag}.parquet")
noise_catalog.to_parquet(noise_file)
noise_file_writer = pq.ParquetWriter(noise_file, SourceCatalog.schema)

Expand Down
2 changes: 2 additions & 0 deletions src/adam_test_data/pointings.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@ class Pointings(qv.Table):
fieldDec_deg = qv.Float64Column()
rotSkyPos_deg = qv.Float64Column()

name = qv.StringAttribute()

# Additional columns which may be useful
observatory_code = qv.LargeStringColumn(nullable=True)

Expand Down
1 change: 1 addition & 0 deletions src/adam_test_data/populations.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ class PhotometricProperties(qv.Table):
class SmallBodies(qv.Table):
orbits = Orbits.as_column()
properties = PhotometricProperties.as_column()
name = qv.StringAttribute()


def photometric_properties_to_sorcha_table(
Expand Down
3 changes: 2 additions & 1 deletion src/adam_test_data/tests/test_main.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ def small_bodies() -> SmallBodies:
H_mf=[18.61],
)

return SmallBodies.from_kwargs(orbits=orbits, properties=properties)
return SmallBodies.from_kwargs(orbits=orbits, properties=properties, name="TestSmallBodies")


@pytest.fixture
Expand Down Expand Up @@ -112,6 +112,7 @@ def pointings() -> Pointings:
],
rotSkyPos_deg=pa.repeat(0.0, num_visits),
observatory_code=pa.repeat("X05", num_visits),
name="TestPointings",
)


Expand Down

0 comments on commit ddc2518

Please sign in to comment.