Skip to content
Snippets Groups Projects

Add inactive sensors to graph

Closed Jorge Prado requested to merge github/fork/RasmusOrsoe/load_inactive_sensors into main
3 files
+ 148
7
Compare changes
  • Side-by-side
  • Inline
Files
3
@@ -3,9 +3,12 @@
from typing import Any, List, Optional, Tuple, Union
import pandas as pd
import sqlite3
import numpy as np
from graphnet.data.dataset import Dataset, ColumnMissingException
from graphnet.data.sqlite.sqlite_utilities import database_table_exists
class SQLiteDataset(Dataset):
"""Pytorch dataset for reading data from SQLite databases."""
@@ -67,11 +70,10 @@ class SQLiteDataset(Dataset):
combined_selections = (
f"{self._index_column} = {index} and {selection}"
)
result = self._conn.execute(
f"SELECT {columns} FROM {table} WHERE "
f"{combined_selections}"
).fetchall()
result = self._conn.execute(
f"SELECT {columns} FROM {table} WHERE "
f"{combined_selections}"
).fetchall()
except sqlite3.OperationalError as e:
if "no such column" in str(e):
raise ColumnMissingException(str(e))
@@ -79,6 +81,23 @@ class SQLiteDataset(Dataset):
raise e
return result
def _get_inactive_sensors(
self,
features: List[Tuple[float, ...]],
columns: List[str],
sequential_index: int,
) -> List[Tuple[float, ...]]:
assert self._conn
index = self._get_event_index(sequential_index)
active_query = f"select (CAST({self._sensor_position['x']} AS str) || '_' || CAST({self._sensor_position['y']} AS str) || '_' || CAST({self._sensor_position['z']} AS str)) as UID from {self._pulsemaps[0]} where {self._index_column} = {index} and {self._selection}"
active_result = self._conn.execute(active_query).fetchall()
columns_str = ", ".join(
columns
) # remove event_no because not in geometry table
query = f"select {columns_str} from {self._geometry_table} where UID not in {str(tuple(np.array(active_result)))}"
inactive_result = self._conn.execute(query).fetchall()
return inactive_result
def _get_all_indices(self) -> List[int]:
self._establish_connection(0)
indices = pd.read_sql_query(
@@ -147,3 +166,15 @@ class SQLiteDataset(Dataset):
self._all_connections_established = False
self._conn = None
return self
def _setup_geometry_table(self, geometry_table: str) -> None:
"""Assign the geometry table to self if it exists in the database."""
assert isinstance(self._path, str)
if database_table_exists(
database_path=self._path, table_name=geometry_table
):
self._geoemtry_table = geometry_table
else:
assert (
1 == 2
), f"Geometry table named {geometry_table} is not in the database {self._path}"
Loading