From ffad792d0d75f0a06e5fac13aa77c780f14352d9 Mon Sep 17 00:00:00 2001 From: Simon Lloyd Date: Thu, 5 Mar 2026 10:54:40 +0000 Subject: [PATCH] refactor(core): re-arange core, move set_default_event and friends out of core. --- .gitignore | 1 + docs/first-steps/data.md | 80 +++---- docs/first-steps/installation.md | 109 +++------ docs/first-steps/workflow.md | 152 ++++++------- docs/usage/api.md | 16 ++ docs/usage/cli.md | 55 +---- docs/usage/defaults.md | 2 +- docs/usage/gui.md | 2 +- docs/usage/index.md | 23 +- docs/usage/shell.md | 1 + docs/usage/tui.md | 1 + pyproject.toml | 1 + src/aimbat/_cli/align.py | 12 +- src/aimbat/_cli/data.py | 95 ++++++-- src/aimbat/_cli/event.py | 213 ++++++++++++++---- src/aimbat/_cli/pick.py | 17 +- src/aimbat/_cli/plot.py | 19 +- src/aimbat/_cli/project.py | 78 ++++++- src/aimbat/_cli/seismogram.py | 152 +++++++++++-- src/aimbat/_cli/snapshot.py | 119 +++++++++- src/aimbat/_cli/station.py | 145 ++++++++++-- src/aimbat/_cli/utils/sampledata.py | 10 +- src/aimbat/_tui/app.py | 144 ++++++++---- src/aimbat/_types/_event.py | 10 +- src/aimbat/_types/_seismogram.py | 4 +- src/aimbat/core/_active_event.py | 8 +- src/aimbat/core/_data.py | 87 ++------ src/aimbat/core/_event.py | 235 ++++++-------------- src/aimbat/core/_iccs.py | 114 +++++++--- src/aimbat/core/_project.py | 116 +++------- src/aimbat/core/_seismogram.py | 206 ++++------------- src/aimbat/core/_snapshot.py | 182 +++------------ src/aimbat/core/_station.py | 218 +++--------------- src/aimbat/models/__init__.py | 1 + src/aimbat/models/_models.py | 5 + src/aimbat/models/_parameters.py | 6 +- src/aimbat/models/_readers.py | 2 + src/aimbat/utils/_style.py | 4 +- tests/integration/core/test_data.py | 99 +-------- tests/integration/core/test_event.py | 105 +++------ tests/integration/core/test_project.py | 28 +-- tests/integration/core/test_seismogram.py | 95 ++------ tests/integration/core/test_snapshots.py | 176 +++------------ tests/integration/core/test_station.py | 79 ++----- tests/integration/models/test_models.py | 10 +- tests/integration/models/test_operations.py | 33 ++- uv.lock | 23 ++ zensical.toml | 9 +- 48 files changed, 1501 insertions(+), 1801 deletions(-) create mode 100644 docs/usage/api.md create mode 100644 docs/usage/shell.md create mode 100644 docs/usage/tui.md diff --git a/.gitignore b/.gitignore index 77124253..5326b7a7 100644 --- a/.gitignore +++ b/.gitignore @@ -38,3 +38,4 @@ aimbat_test*.log GEMINI.md CLAUDE.md .claude/settings.local.json +scratch/ diff --git a/docs/first-steps/data.md b/docs/first-steps/data.md index 472a0701..a84d1f45 100644 --- a/docs/first-steps/data.md +++ b/docs/first-steps/data.md @@ -1,24 +1,17 @@ # Data -Because a lot of operations in AIMBAT involve adjusting parameters that do -*not* require reading entire seismograms, the time series components of the -seismograms are kept separate from the rest of the data. If this were not the -case, doing something like adjusting the time window for the cross-correlation -for an event with 200 seismograms, would require reading all 200 seismograms -from disk, even though their data are not actually used for that operation. Not -having to do this naturally benefits performance, but also means that an AIMBAT -project does not really use up much disk space at all. As this is true for any -amount of seismograms, we can include multiple events in a single project. +AIMBAT treats input data as read-only. Processing parameters and results are +stored separately from the data in a database. Once imported into a project, +the data sources (e.g. SAC files) are only used to read the waveforms. All +metadata (e.g. event and station information) is stored in the database. ## Data hierarchy Moving away from storing all data in individual files means things are -organised differently. Most importantly, seismograms, after importing them into -a project, no longer each contain event and station information. Instead they -use (and share) separately stored events and stations. Thus a single event will -contain multiple seismograms. The same is true for a single station, while the -seismograms themselves use exactly one event, one station, and one data source -(containing the time series). +organised differently. For example, a seismogram in AIMBAT is no longer a file; +it is an object in the database that links to a data source, station and event. +Stations and events are also objects in the database, and they are shared +between seismograms. ```mermaid --- @@ -32,21 +25,25 @@ erDiagram ``` -While all of this happens transparently to the user, there are still some -things to be mindful of: +!!! Tip + Seismograms that are supposed to be processed together are identified only + by their shared event and station in the database. You can therefore be + quite flexible in how you organise your data on disk, but you must make + sure they reference the exact same station and event information — small + differences in the metadata (e.g. a different number of decimal places) may + lead to AIMBAT treating seismograms as belonging to different events or + stations. + +## Deleting items + +All of the above happens transparently to the user, but there are some things to +be mindful of when it comes to deleting items from a project: -- Seismograms that are supposed to be processed together are no longer - identified e.g. by being stored in the same directory, but instead from using - the same event and station. Thus the files they are created from when adding - them to an AIMBAT project *must* contain identical station and event data. - Deleting[^1] an event or station from a project will remove all related seismograms. -- Conversely, deleting a seismogram will *not* remove the event or - station. This remains true even if an event or station end up not being used - by any seismograms in the project. -- After importing a seismogram into a project, it no longer accesses the - metadata stored in the source file. It does, however, still need to be able - to read the time series data from it. +- Conversely, deleting a seismogram will *not* remove the event or station. + This remains true even if an event or station end up not being used by any + seismograms in the project. !!! Tip If this all seems overly complicated to you, remember that often parameters @@ -88,18 +85,13 @@ processing: ### AIMBAT Defaults -AIMBAT defaults are global settings that control how the application itself -behaves, as well as defaults for [event](#event-parameters) and -[seismogram](#seismogram-parameters) parameters when they are instantiated. -The currently used values for these parameters can be found by running -`#!bash aimbat settings` in your terminal. As some settings are relevant before -a project is created, they cannot be stored in the project file. To override these -settings you can set the corresponding environment variable directly (e.g. -`export AIMBAT_PROJECT=different_project_name.db`) or place those settings in a -`.env`[^2] file. Note that if you set them in both places the environment -variable is used. - -[^2]: The file must be named exactly `.env`, and not `SOMENAME.env`! +AIMBAT [defaults](../usage/defaults.md) are global settings that control how +the application itself behaves, as well as defaults for +[event](#event-parameters) and [seismogram](#seismogram-parameters) parameters +when they are instantiated. The currently used values for these parameters can +be listed by running `#!bash aimbat utils settings` in the terminal. As some +settings are relevant before a project is created, they cannot be stored in +the project file. ### Event Parameters @@ -108,14 +100,14 @@ specific to an event (e.g. if an event should be marked as completed), or parameters that are shared across all seismograms of that event (e.g. time window for the cross-correlation, filter parameters, etc.). These parameters are attributes of the -[`AimbatEventParametersBase`][aimbat.lib.models.AimbatEventParametersBase] +[`AimbatEventParametersBase`][aimbat.models.AimbatEventParametersBase] class. ### Seismogram Parameters Seismogram parameters are also used during processing. Most notably the time picks belong to this tier. These parameters are attributes of the -[`AimbatSeismogramParametersBase`][aimbat.lib.models.AimbatSeismogramParametersBase] +[`AimbatSeismogramParametersBase`][aimbat.models.AimbatSeismogramParametersBase] class. ## Snapshots @@ -138,7 +130,7 @@ Internally, the items in a project use [UUIDs](https://en.wikipedia.org/wiki/Universally_unique_identifier) to identify themselves. They look something like this: -``` +```text 37a8245f-c508-46a7-9bbc-d1c601e42983 ``` @@ -149,7 +141,7 @@ easier, they are typically presented in a truncated form to the user (but always long enough to be unique). For example, if there are only four seismograms and they have these UUIDS: -``` +```text 6a4acdf7-6c7b-4523-aaaa-0a674cdc5f2d 647568aa-8361-45ef-bfc8-61f873847f17 c980918d-106d-44d9-a3fa-5740f58edf4e @@ -158,7 +150,7 @@ c980918d-106d-44d9-a3fa-5740f58edf4e they can still be unambiguously identified using only the first two characters: -``` +```text 6a 64 c9 diff --git a/docs/first-steps/installation.md b/docs/first-steps/installation.md index aa0574ab..9b977aaa 100644 --- a/docs/first-steps/installation.md +++ b/docs/first-steps/installation.md @@ -1,18 +1,19 @@ # Installing AIMBAT -AIMBAT is built on top of standard [Python](https://www.python.org) and uses -some popular third party modules (e.g. [NumPy][numpy], [SciPy][scipy]). In -order to benefit from modern Python features and up to date modules, AIMBAT is -developed on the latest stable Python versions. Automatic tests are done on -version 3.12 and newer. - -AIMBAT is available as a package from the -[Python Package Index](https://pypi.org/project/aimbat/). This means it can be -installed using the [`pip`](https://pip.pypa.io/en/stable/) module. However, as -AIMBAT is a standalone application (rather than a library), we recommend -installing it using [`uv`](https://docs.astral.sh/uv/) instead. `uv` is a -single binary that doesn't require any dependencies to be installed, and it -allows you to install and run AIMBAT in an isolated environment. +AIMBAT is built on top of standard [Python](https://www.python.org), combined +with a number of popular third party modules such as [NumPy][numpy] and +[SciPy][scipy]. It is available as a package from the +[Python Package Index][pypi], and can therefore be +installed using Python's standard package manager, +[`pip`](https://pip.pypa.io/en/stable/). + +However, as AIMBAT is a standalone application, we recommend installing it +using tools like [`uv`][uv] or [`pipx`][pipx] instead[^1]. Here we show how to +install and run AIMBAT using `uv`. + +[^1]: If you want to get really fancy you can also use tools like + [MISE-EN-PLACE](https://mise.jdx.dev), which use `uv` or `pipx` under the + hood. ## Running AIMBAT without installing @@ -22,44 +23,21 @@ of using `uv` is that it can run applications without installing them: ```bash $ # First check that uv is available: -$ uv --version -uv 0.8.14 +$ uv self version +uv 0.10.6 $ # Next run AIMBAT using uv tool: -$ uv tool run aimbat -⠦ Resolving dependencies... -... -Usage: aimbat COMMAND - -AIMBAT command line interface entrypoint for all other commands. - -This is the main command line interface for AIMBAT. It must be -executed with a command (as specified below) to actually do anything. -Help for individual commands is available by typing aimbat COMMAND ---help. - -╭─ Commands ────────────────────────────────────────────────────────╮ -│ data Manage seismogram files in an AIMBAT project. │ -│ event View and manage events in the AIMBAT project. │ -│ iccs ICCS processing tools. │ -│ project Manage AIMBAT projects. │ -│ seismogram View and manage seismograms in the AIMBAT project. │ -│ settings Print a table with default settings used in AIMBAT. │ -│ snapshot View and manage snapshots. │ -│ station View and manage stations. │ -│ utils Utilities for AIMBAT. │ -│ --help -h Display this message and exit. │ -│ --version Display application version. │ -╰───────────────────────────────────────────────────────────────────╯ +$ uv tool run aimbat --version +⠴ Resolving dependencies... +2.0.0 ``` -This will likely have taken a while, as `uv` has to download and install a lot -of dependencies. Subsequent runs will be much faster. `uv tool run` is such a +The initial run may take a while to download AIMBAT and its dependencies, but +once that is done subsequent runs will be much faster. `uv tool run` is such a useful command that it has its own alias, `uvx`: ```bash -# 'uvx' is an alias for 'uv tool run': -$ uvx aimbat -Usage: aimbat COMMAND +$ uvx aimbat --version +2.0.0 ... ``` @@ -82,17 +60,16 @@ To clean up after yourself, you can remove the `uv` cache: ```bash $ uv clean -Clearing cache at: /home/USER/.cache/uv +Clearing cache at: /home/bob/.cache/uv +Cleaning [==================> ] 91% Removed 26702 files (2.1GiB) ``` -## Installing AIMBAT locally +## Installing locally -After successfully test driving AIMBAT you may want to actually install it. -This is very simple with `uv`: +Permanent installation is just as easy. Just run `uv tool install`: ```bash -# Again use uv tool, but this time with the 'install' command: $ uv tool install aimbat ⠦ Resolving dependencies... ⠇ Preparing packages... (66/69) @@ -110,12 +87,12 @@ Usage: aimbat COMMAND ``` !!! tip - If the above command fails (because your shell can't find the `aimbat` command), you may need to add `~/.local/bin` to your `PATH`. This can be done automatically by running `#!bash uv tool update-shell`. -Upgrading or uninstalling AIMBAT is just as easy: +Upgrading or uninstalling is just as easy with `uv tool upgrade` and `uv tool +uninstall`: ```bash $ uv tool upgrade aimbat @@ -124,26 +101,8 @@ $ uv tool uninstall aimbat Uninstalled 1 executable: aimbat ``` -## Demo - -This demo shows what you can expect to see in your terminal when running -the commands above: - -
- - - -!!! tip - - These kinds of recordings work just like normal videos, but you can also - select text and copy it to your clipboard! +[numpy]: https://numpy.org +[scipy]: https://scipy.org +[pypi]: https://pypi.org/project/aimbat/ +[uv]: https://docs.astral.sh/uv/ +[pipx]: https://pipx.pypa.io diff --git a/docs/first-steps/workflow.md b/docs/first-steps/workflow.md index d86310cf..9f92252a 100644 --- a/docs/first-steps/workflow.md +++ b/docs/first-steps/workflow.md @@ -1,12 +1,26 @@ # Workflow and Strategy +The best results with AIMBAT are obtained if you understand the unique, +non-linear workflow that it prescribes. Unlike more traditional top-down +processing, AIMBAT is more of an iterative process whereby parameters are +constantly adjusted to gradually determine the optimal settings. The exact +strategy for a particular event may differ, but there are some general +guidelines below that we recommend following. + ## Without AIMBAT -Multi-Channel Cross-Correlation[@vandecar_determination_1990] (MCCC) relies on -narrow time windows, focused on the initial arrival arrival of the targeted +Multi-Channel Cross-Correlation[^1] (MCCC) relies on +narrow time windows, focused on the initial arrival of the targeted phase in order to yield high quality results. Manually picking the phase -arrival on each seismogram individually is a very time consuming task, which we -highlight with stacked cards in the below flowchart: +arrival on each seismogram individually is a very time consuming task, which is +highlighted by the stacked cards in the flowchart below: + +[^1]: + VanDecar, J. C., and R. S. Crosson. “Determination of Teleseismic + Relative Phase Arrival Times Using Multi-Channel Cross-Correlation and + Least Squares.” Bulletin of the Seismological Society of America, + vol. 80, no. 1, Feb. 1990, pp. 150–69, + . ```mermaid flowchart TD @@ -20,19 +34,24 @@ flowchart TD ## With AIMBAT -AIMBAT[@lou_aimbat_2013] stacks all input seismograms (aligned on the picked +AIMBAT[^2] stacks all input seismograms (aligned on the picked phase arrival) and operates on that stack instead of individual seismograms. This allows picking the phase arrival once for all seismograms simultaneously, and then improving it iteratively before running MCCC. Note that both the ICCS algorithm, as well as adjusting AIMBAT parameters are iterative processes. +[^2]: + Lou, X., et al. “AIMBAT: A Python/Matplotlib Tool for Measuring + Teleseismic Arrival Times.” Seismological Research Letters, vol. 84, + no. 1, Jan. 2013, pp. 85–93, . + ``` mermaid flowchart TD A@{ shape: circle, label: "Start"} A --> B>Import seismograms containing initial picks t0]; B --> F E[Adjust AIMBAT parameters]; - E --> F>Run ICCS with intial/updated parameters] + E --> F>Run ICCS with initial/updated parameters] F --> G[Inspect results of alignment]; G --> H{"Continue with @@ -43,86 +62,41 @@ flowchart TD ## Strategy -AIMBAT does not prescribe a single strategy for picking processing parameters. -Generally speaking, we recommend adjusting only one parameter at a time between -ICCS runs, and prioritising them as follows: - -1. Filter parameters. -2. Selection of high quality seismograms. -3. Time window boundaries -4. Manually picking phase arrival. - -!!! tip - - Remember that you can create snapshots of the current AIMBAT parameters at - any time, and then rollback to that state if you notice you went into the - wrong direction. We therefore encourage experimenting a bit with the - strategy, as different events may require doing things slightly - differently. - -``` mermaid ---- -title: AIMBAT Workflow ---- -flowchart TD - A[Start] --> B>Check data]; - B --> C{"Any - errors?"}; - C --->|No| G>Import files to AIMBAT and - run ICCS with initial picks - and default parameters]; - C -->|Yes| F[Fix files]; - F --> B; - G --> I["Inspect initial results"]; - - I --> Iq2{"Adjust - filtering?"}; - Iq2 -->|No| Iq3{"Any bad - traces?"}; - Iq3 -->|No| Iq4{"Adjust time - window?"}; - Iq4 -->|No| Iq5{"Has the phase - arrival emerged - in stack?"}; - Iq5 -->|No| Irerun; - - Iq2 -->|Yes| Iq2y["Set new filter parameters."]; - Iq2y --> Iq2yq{"Re-run - ICCS - now?"}; - Iq2yq -->|No|Iq3; - Iq2yq -->|Yes|Irerun; - - Iq3 -->|Yes| Iq3y["Select/deselect seismograms."]; - Iq3y --> Iq3yq{"Re-run - ICCS - now?"}; - Iq3yq -->|No|Iq4; - Iq3yq -->|Yes|Irerun; - - Iq4 -->|Yes| Iq4y["Pick new time window"]; - Iq4y --> Iq4yq{"Re-run - ICCS - now?"}; - Iq4yq -->|No|Iq5; - Iq4yq -->|Yes|Irerun; - - Iq5 -->|Yes| Iq5q{"Is the pick - on the visible - arrival?"}; - Iq5q -->|No| Iq5qy["Pick new Time"] --> Irerun; - Iq5q -->|Yes| Irerun; - - Irerun>"Run ICCS with - updated settings"] --> I2; - - I2["Inspect updated results"]; - - I2 --> qM{"Continue - with - MCCC?"} - qM -->|Yes| M>"MCCC with final pick - and time window"] --> Z[END]; - qM -->|No| Iq2; - -``` +AIMBAT does not prescribe a single strategy for picking/updating processing +parameters. That said, some general principles to follow are: + +1. Only change one parameter at a time, and then run ICCS to see the effect + of that change on the alignment +2. Snapshots are immediate and use no storage, so create them often (and don't + skip adding a comment that describes the snapshot). +3. Don't get too distracted by individual seismograms that are not well aligned + or seem of poor quality. Trust the algorithm to deal with those. + +### ICCS running modes + +Points 1 and 2 above are pretty self-explanatory, but point 3 deserves a bit +more explanation. The ICCS algorithm has two flags that can be set before each +run: + +- **Autoflip**: If set to `True`, the algorithm will toggle the ``flip`` + parameter for seismograms that are negatively correlated with the stack (i.e. + the maximum absolute correlation coefficient is negative). +- **Autoselect**: If set to `True`, the algorithm will automatically set the + `select` parameter for seismograms that are poorly correlated with the stack + to `False`, but it will also toggle previously deselected seismograms back + to `True` if they become well correlated with the stack. + +The `flip` parameter determines whether a seismogram is flipped in polarity as +the data are prepared for the stack and following cross-correlation. The +`select` parameter determines whether a seismogram is included in the stack; +however, *all* seismograms are cross-correlated with that stack regardless of +their `select` status. This means that all seismograms, even with +`select=False`, can self-recover if they fit the stack better after changing +parameters. + +!!! Tip + If seismograms are so bad that they wander off into the distance (i.e the + difference between initial pick and revised pick after running ICCS is + very large), consider deleting them completely from the AIMBAT project. + This will prevent these rogue seismograms from influencing the valid + ranges for updating time windows and picks. diff --git a/docs/usage/api.md b/docs/usage/api.md new file mode 100644 index 00000000..a3394e07 --- /dev/null +++ b/docs/usage/api.md @@ -0,0 +1,16 @@ +# Python API + +If none of the above interfaces suit your needs, or you want to write custom +scripts, you can use the AIMBAT Python API. This is the most powerful way to +interact with your projects. + +## Core Concepts + +The API is built on three main components: + +1. **Models**: [SQLModel](https://sqlmodel.tiangolo.com) classes that represent + the database schema (`aimbat.models`) as Python objects. +2. **Core Functions**: High-level operations that manipulate those models + (`aimbat.core`). +3. **Database Session**: A SQLAlchemy session used to track changes and + interact with the project database. diff --git a/docs/usage/cli.md b/docs/usage/cli.md index 8dec31c0..eec61e8c 100644 --- a/docs/usage/cli.md +++ b/docs/usage/cli.md @@ -1,43 +1,12 @@ -# Command Line - -## Seismogram files - -AIMBAT uses [SAC](https://ds.iris.edu/files/sac-manual/) files as input. Before -adding files to an AIMBAT project please ensure the following header fields are -set correctly in all files: - -- Seismogram begin time (SAC header *B*). -- Seismogram reference time (*KZTIME*) and date (*KZDATE*). -- Station name (*KSTNM*), latitude (*STLA*) and longitude *STLO*). -- Event origin time (*O*), latitude (*EVLA*) and longitude (*EVLO*). - -To detect any potential problems with the data before importing them into AIMBAT, -you can use the AIMBAT cli: - - - -```bash -$ aimbat checkdata *.BHZ -sacfile_01.BHZ: ✓✓✓ -sacfile_02.BHZ: ✓✓✓ -sacfile_03.BHZ: ✓✓✓ -sacfile_04.BHZ: ✓✓✓ -sacfile_05.BHZ: ✓✓✓ -sacfile_06.BHZ: ✓✓✓ -sacfile_07.BHZ: ✓✓✓ -sacfile_08.BHZ: ✓✓✓ -sacfile_09.BHZ: ✓✓✓ -sacfile_10.BHZ: ✓✓✓ -sacfile_11.BHZ: ✓✓✓ -... -sacfile_NN.BHZ: ✓✓✓ - -No issues found! -``` - -The seismogram files can be stored in an arbitrary directory (i.e. they do not -necessarily need to be stored together with an AIMBAT project file). - -!!! warning - After importing files into an AIMBAT project, their location (and contents) - should not be changed! +# Command Line Interface (CLI) + +The AIMBAT CLI is the primary tool for administrative tasks like creating projects, +importing data, and managing snapshots. It is also suitable for batch +processing and scripting. + +!!! Warning "Parameter Validation" + The CLI performs basic validation of processing parameters (e.g., ensuring + values are provided in the correct format), but it does *not* perform the + same data-aware validation as the other interfaces. For example, the TUI + will prevent you from setting a time window that extends beyond the + available data for an event. diff --git a/docs/usage/defaults.md b/docs/usage/defaults.md index dd0a64c7..da5a529a 100644 --- a/docs/usage/defaults.md +++ b/docs/usage/defaults.md @@ -1,4 +1,4 @@ -# Defaults +# AIMBAT defaults AIMBAT behaviour can be customised via the following settings. Each setting can be overridden on a per-project basis (in order of precedence): diff --git a/docs/usage/gui.md b/docs/usage/gui.md index 150f5491..9bb1e1c3 100644 --- a/docs/usage/gui.md +++ b/docs/usage/gui.md @@ -1 +1 @@ -# Graphical UI +# Graphical User Interface (GUI) diff --git a/docs/usage/index.md b/docs/usage/index.md index a1db4314..fca7f2be 100644 --- a/docs/usage/index.md +++ b/docs/usage/index.md @@ -1,13 +1,20 @@ # Using AIMBAT -Once [installed](../first-steps/installation.md), AIMBAT can be used in several ways: +Once [installed](../first-steps/installation.md), AIMBAT can be used in several +ways, each of which has unique strengths depending on the task at hand: -- As a [command line](cli.md) application. -- As a [GUI](gui.md) (graphical user interface) application. +- **[Command line](cli.md)**: Ideal for administrative tasks like + adding data to a project, and exploring the data after they are added. +- **[Interactive Shell](shell.md)**: Similar to the CLI but with the added benefit of extra + context and command history. Unlike the CLI, it is not possible to set + parameters to nonsensical values. +- **[Terminal UI](tui.md)**: This is where processing typically happens. + The TUI is designed for efficient, mouse-free, and keyboard-driven navigation + for users familiar with the workflow. +- **[Graphical UI](gui.md)**: Mouse driven interface for users who prefer a + more visual approach. The GUI is ideal for newer users who need more guidance + and visual cues to navigate the workflow. +- **[Python API](api.md)**: The preferred way for scripting and automated +processing by writing custom Python scripts. Complete walkthroughs for each of these options are presented in the following sections. - -!!! tip - - AIMBAT project files are independent of user interface. In other words, - existing projects can be opened in any way (cli or gui). diff --git a/docs/usage/shell.md b/docs/usage/shell.md new file mode 100644 index 00000000..5b0c43c8 --- /dev/null +++ b/docs/usage/shell.md @@ -0,0 +1 @@ +# AIMBAT Shell diff --git a/docs/usage/tui.md b/docs/usage/tui.md new file mode 100644 index 00000000..202edba3 --- /dev/null +++ b/docs/usage/tui.md @@ -0,0 +1 @@ +# Terminal User Interface (TUI) diff --git a/pyproject.toml b/pyproject.toml index 73cd6e81..994fcd5a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ "pandas-stubs>=3.0.0.260204", "textual>=8.0.0", "textual-fspicker>=1.0.0", + "prompt-toolkit>=3.0.52", ] [project.urls] diff --git a/src/aimbat/_cli/align.py b/src/aimbat/_cli/align.py index 31f5c5e1..ff0f9432 100644 --- a/src/aimbat/_cli/align.py +++ b/src/aimbat/_cli/align.py @@ -33,11 +33,12 @@ def cli_iccs_run( cross-correlation with the stack falls below `min_ccnorm`. """ from aimbat.db import engine - from aimbat.core import create_iccs_instance, run_iccs + from aimbat.core import create_iccs_instance, run_iccs, get_active_event from sqlmodel import Session with Session(engine) as session: - iccs = create_iccs_instance(session) + active_event = get_active_event(session) + iccs = create_iccs_instance(session, active_event).iccs run_iccs(session, iccs, autoflip, autoselect) @@ -58,12 +59,13 @@ def cli_mccc_run( the currently selected ones. """ from aimbat.db import engine - from aimbat.core import create_iccs_instance, run_mccc + from aimbat.core import create_iccs_instance, run_mccc, get_active_event from sqlmodel import Session with Session(engine) as session: - iccs = create_iccs_instance(session) - run_mccc(session, iccs, all_seismograms) + active_event = get_active_event(session) + iccs = create_iccs_instance(session, active_event).iccs + run_mccc(session, active_event, iccs, all_seismograms) if __name__ == "__main__": diff --git a/src/aimbat/_cli/data.py b/src/aimbat/_cli/data.py index 723b32de..0dae4035 100644 --- a/src/aimbat/_cli/data.py +++ b/src/aimbat/_cli/data.py @@ -27,6 +27,7 @@ records are reused rather than duplicated. """ +import uuid from .common import ( GlobalParameters, TableParameters, @@ -35,13 +36,12 @@ use_station_parameter, use_event_parameter, ) -from aimbat.models import AimbatEvent, AimbatStation +from aimbat.models import AimbatEvent, AimbatStation, AimbatDataSource from aimbat.io import DataType -from sqlmodel import Session +from sqlmodel import Session, select from cyclopts import App, Parameter, validators from pathlib import Path from typing import Annotated -import uuid app = App(name="data", help=__doc__, help_format="markdown") @@ -106,22 +106,6 @@ def cli_data_add( ) -@app.command(name="list") -@simple_exception -def cli_data_list( - *, - all_events: Annotated[bool, ALL_EVENTS_PARAMETER] = False, - table_parameters: TableParameters = TableParameters(), - global_parameters: GlobalParameters = GlobalParameters(), -) -> None: - """Print a table of data sources registered in the AIMBAT project.""" - from aimbat.db import engine - from aimbat.core import print_data_table - - with Session(engine) as session: - print_data_table(session, table_parameters.short, all_events) - - @app.command(name="dump") @simple_exception def cli_data_dump( @@ -140,5 +124,78 @@ def cli_data_dump( print_json(dump_data_table_to_json(session)) +@app.command(name="list") +@simple_exception +def cli_data_list( + *, + all_events: Annotated[bool, ALL_EVENTS_PARAMETER] = False, + table_parameters: TableParameters = TableParameters(), + global_parameters: GlobalParameters = GlobalParameters(), +) -> None: + """Print a table of data sources registered in the AIMBAT project.""" + from aimbat.db import engine + from aimbat.core import get_active_event, get_data_for_event + from aimbat.utils import uuid_shortener, make_table, TABLE_STYLING + from aimbat.logger import logger + from rich.console import Console + + short = table_parameters.short + + with Session(engine) as session: + logger.info("Printing data sources table.") + + if all_events: + aimbat_data_sources = session.exec(select(AimbatDataSource)).all() + title = "Data sources for all events" + else: + active_event = get_active_event(session) + aimbat_data_sources = get_data_for_event(session, active_event) + time = ( + active_event.time.strftime("%Y-%m-%d %H:%M:%S") + if short + else active_event.time + ) + id = uuid_shortener(session, active_event) if short else active_event.id + title = f"Data sources for event {time} (ID={id})" + + logger.debug(f"Found {len(aimbat_data_sources)} data sources in total.") + + rows = [ + [ + uuid_shortener(session, a) if short else str(a.id), + str(a.datatype), + str(a.sourcename), + ( + uuid_shortener(session, a.seismogram) + if short + else str(a.seismogram.id) + ), + ] + for a in aimbat_data_sources + ] + + table = make_table(title=title) + + table.add_column( + "ID (shortened)" if short else "ID", + justify="center", + style=TABLE_STYLING.id, + no_wrap=True, + ) + table.add_column("Datatype", justify="center", style=TABLE_STYLING.mine) + table.add_column( + "Source", justify="left", style=TABLE_STYLING.mine, no_wrap=True + ) + table.add_column( + "Seismogram ID", justify="center", style=TABLE_STYLING.linked, no_wrap=True + ) + + for row in rows: + table.add_row(*row) + + console = Console() + console.print(table) + + if __name__ == "__main__": app() diff --git a/src/aimbat/_cli/event.py b/src/aimbat/_cli/event.py index 66ba82b4..42ba4dea 100644 --- a/src/aimbat/_cli/event.py +++ b/src/aimbat/_cli/event.py @@ -52,43 +52,6 @@ def cli_event_activate( set_active_event_by_id(session, event_id) -@app.command(name="dump") -@simple_exception -def cli_event_dump( - *, - global_parameters: GlobalParameters = GlobalParameters(), -) -> None: - """Dump the contents of the AIMBAT event table to JSON. - - Output can be piped or redirected for use in external tools or scripts. - """ - from aimbat.db import engine - from aimbat.core import dump_event_table_to_json - from rich import print_json - - with Session(engine) as session: - print_json(dump_event_table_to_json(session)) - - -@app.command(name="list") -@simple_exception -def cli_event_list( - *, - table_parameters: TableParameters = TableParameters(), - global_parameters: GlobalParameters = GlobalParameters(), -) -> None: - """Print a table of events stored in the AIMBAT project. - - The active event is highlighted. Use `event activate` to change which event - is processed by subsequent commands. - """ - from aimbat.db import engine - from aimbat.core import print_event_table - - with Session(engine) as session: - print_event_table(session, table_parameters.short) - - @parameter.command(name="get") @simple_exception def cli_event_parameter_get( @@ -103,11 +66,12 @@ def cli_event_parameter_get( """ from aimbat.db import engine - from aimbat.core import get_event_parameter + from aimbat.core import get_event_parameter, get_active_event from sqlmodel import Session with Session(engine) as session: - value = get_event_parameter(session, name) + active_event = get_active_event(session) + value = get_event_parameter(session, active_event, name) if isinstance(value, Timedelta): print(f"{value.total_seconds()}s") else: @@ -130,7 +94,7 @@ def cli_event_parameter_set( are interpreted as seconds. """ from aimbat.db import engine - from aimbat.core import set_event_parameter + from aimbat.core import set_event_parameter, get_active_event from sqlmodel import Session _TIMEDELTA_PARAMS = (EventParameter.WINDOW_PRE, EventParameter.WINDOW_POST) @@ -142,7 +106,8 @@ def cli_event_parameter_set( parsed_value = Timedelta(value) with Session(engine) as session: - set_event_parameter(session, name, parsed_value) + active_event = get_active_event(session) + set_event_parameter(session, active_event, name, parsed_value) @parameter.command(name="dump") @@ -153,13 +118,121 @@ def cli_event_parameter_dump( ) -> None: """Dump event parameter table to json.""" from aimbat.db import engine - from aimbat.core import dump_event_parameter_table_to_json + from aimbat.core import dump_event_parameter_table_to_json, get_active_event from sqlmodel import Session from rich import print_json with Session(engine) as session: + active_event = get_active_event(session) if not all_events else None print_json( - dump_event_parameter_table_to_json(session, all_events, as_string=True) + dump_event_parameter_table_to_json( + session, all_events, as_string=True, event=active_event + ) + ) + + +@app.command(name="dump") +@simple_exception +def cli_event_dump( + *, + global_parameters: GlobalParameters = GlobalParameters(), +) -> None: + """Dump the contents of the AIMBAT event table to JSON. + + Output can be piped or redirected for use in external tools or scripts. + """ + from aimbat.db import engine + from aimbat.core import dump_event_table_to_json + from rich import print_json + + with Session(engine) as session: + print_json(dump_event_table_to_json(session)) + + +@app.command(name="list") +@simple_exception +def cli_event_list( + *, + table_parameters: TableParameters = TableParameters(), + global_parameters: GlobalParameters = GlobalParameters(), +) -> None: + """Print a table of events stored in the AIMBAT project. + + The active event is highlighted. Use `event activate` to change which event + is processed by subsequent commands. + """ + from aimbat.db import engine + from aimbat.core import dump_event_table_to_json + from aimbat.utils import uuid_shortener, json_to_table, TABLE_STYLING + from aimbat.logger import logger + from pandas import Timestamp + + short = table_parameters.short + + with Session(engine) as session: + logger.info("Printing AIMBAT events table.") + + json_to_table( + data=dump_event_table_to_json(session, as_string=False), + title="AIMBAT Events", + column_order=[ + "id", + "active", + "time", + "latitude", + "longitude", + "depth", + "completed", + "seismogram_count", + "station_count", + ], + formatters={ + "id": lambda x: ( + uuid_shortener(session, AimbatEvent, str_uuid=x) if short else x + ), + "active": TABLE_STYLING.bool_formatter, + "time": lambda x: TABLE_STYLING.timestamp_formatter( + Timestamp(x), short + ), + "latitude": lambda x: f"{x:.3f}" if short else str(x), + "longitude": lambda x: f"{x:.3f}" if short else str(x), + "depth": lambda x: f"{x:.0f}" if short and x is not None else str(x), + "completed": TABLE_STYLING.bool_formatter, + "last_modified": lambda x: TABLE_STYLING.timestamp_formatter( + Timestamp(x), short + ), + }, + common_column_kwargs={"justify": "center"}, + column_kwargs={ + "id": { + "header": "ID (shortened)" if short else "ID", + "style": TABLE_STYLING.id, + "no_wrap": True, + }, + "active": {"style": TABLE_STYLING.mine, "no_wrap": True}, + "time": { + "header": "Date & Time", + "style": TABLE_STYLING.mine, + "no_wrap": True, + }, + "last_modified": { + "header": "Last Modified", + "style": TABLE_STYLING.mine, + "no_wrap": True, + }, + "latitude": {"style": TABLE_STYLING.mine}, + "longitude": {"style": TABLE_STYLING.mine}, + "depth": {"style": TABLE_STYLING.mine}, + "completed": {"style": TABLE_STYLING.parameters}, + "seismogram_count": { + "header": "# Seismograms", + "style": TABLE_STYLING.linked, + }, + "station_count": { + "header": "# Stations", + "style": TABLE_STYLING.linked, + }, + }, ) @@ -177,11 +250,59 @@ def cli_event_parameter_list( """ from aimbat.db import engine - from aimbat.core import print_event_parameter_table - from sqlmodel import Session + from aimbat.core import dump_event_parameter_table_to_json, get_active_event + from aimbat.utils import uuid_shortener, json_to_table, TABLE_STYLING + from aimbat.logger import logger + + short = table_parameters.short with Session(engine) as session: - print_event_parameter_table(session, table_parameters.short, all_events) + if all_events: + logger.info("Printing AIMBAT event parameters table for all events.") + json_to_table( + data=dump_event_parameter_table_to_json( + session, all_events=True, as_string=False + ), + title="Event parameters for all events", + skip_keys=["id"], + column_order=[ + "event_id", + "completed", + "window_pre", + "window_post", + "min_ccnorm", + ], + formatters={ + "event_id": lambda x: ( + uuid_shortener(session, AimbatEvent, str_uuid=x) if short else x + ), + }, + common_column_kwargs={"highlight": True}, + column_kwargs={ + "event_id": { + "header": "Event ID (shortened)" if short else "Event ID", + "justify": "center", + "style": TABLE_STYLING.mine, + }, + }, + ) + else: + logger.info("Printing AIMBAT event parameters table for active event.") + + active_event = get_active_event(session) + json_to_table( + data=active_event.parameters.model_dump(mode="json"), + title=f"Event parameters for event: {uuid_shortener(session, active_event) if short else str(active_event.id)}", + skip_keys=["id", "event_id"], + common_column_kwargs={"highlight": True}, + column_kwargs={ + "Key": { + "header": "Parameter", + "justify": "left", + "style": TABLE_STYLING.id, + }, + }, + ) if __name__ == "__main__": diff --git a/src/aimbat/_cli/pick.py b/src/aimbat/_cli/pick.py index 05b57ac0..8e2c750e 100644 --- a/src/aimbat/_cli/pick.py +++ b/src/aimbat/_cli/pick.py @@ -30,11 +30,12 @@ def cli_update_phase_pick( use_seismogram_image: Use the seismogram image to update pick. """ from aimbat.db import engine - from aimbat.core import create_iccs_instance, update_pick + from aimbat.core import create_iccs_instance, update_pick, get_active_event from sqlmodel import Session with Session(engine) as session: - iccs = create_iccs_instance(session) + active_event = get_active_event(session) + iccs = create_iccs_instance(session, active_event).iccs update_pick( session, iccs, @@ -63,13 +64,15 @@ def cli_pick_timewindow( use_seismogram_image: Use the seismogram image to pick the time window. """ from aimbat.db import engine - from aimbat.core import create_iccs_instance, update_timewindow + from aimbat.core import create_iccs_instance, update_timewindow, get_active_event from sqlmodel import Session with Session(engine) as session: - iccs = create_iccs_instance(session) + active_event = get_active_event(session) + iccs = create_iccs_instance(session, active_event).iccs update_timewindow( session, + active_event, iccs, iccs_parameters.context, iccs_parameters.all, @@ -92,13 +95,15 @@ def cli_pick_min_ccnorm( automatically de-selected when running ICCS with `--autoselect`. """ from aimbat.db import engine - from aimbat.core import create_iccs_instance, update_min_ccnorm + from aimbat.core import create_iccs_instance, update_min_ccnorm, get_active_event from sqlmodel import Session with Session(engine) as session: - iccs = create_iccs_instance(session) + active_event = get_active_event(session) + iccs = create_iccs_instance(session, active_event).iccs update_min_ccnorm( session, + active_event, iccs, iccs_parameters.context, iccs_parameters.all, diff --git a/src/aimbat/_cli/plot.py b/src/aimbat/_cli/plot.py index 7c50e3c4..8da06215 100644 --- a/src/aimbat/_cli/plot.py +++ b/src/aimbat/_cli/plot.py @@ -28,11 +28,12 @@ def cli_seismogram_plot( ) -> None: """Plot raw seismograms for the active event sorted by epicentral distance.""" from aimbat.db import engine - from aimbat.core import plot_all_seismograms + from aimbat.core import plot_all_seismograms, get_active_event from sqlmodel import Session with Session(engine) as session: - plot_all_seismograms(session, return_fig=False) + active_event = get_active_event(session) + plot_all_seismograms(session, active_event, return_fig=False) @app.command(name="stack") @@ -44,11 +45,12 @@ def cli_iccs_plot_stack( ) -> None: """Plot the ICCS stack of the active event.""" from aimbat.db import engine - from aimbat.core import create_iccs_instance, plot_stack + from aimbat.core import create_iccs_instance, plot_stack, get_active_event from sqlmodel import Session with Session(engine) as session: - iccs = create_iccs_instance(session) + active_event = get_active_event(session) + iccs = create_iccs_instance(session, active_event).iccs plot_stack(iccs, iccs_parameters.context, iccs_parameters.all, return_fig=False) @@ -61,11 +63,16 @@ def cli_iccs_plot_image( ) -> None: """Plot the ICCS seismograms of the active event as an image.""" from aimbat.db import engine - from aimbat.core import create_iccs_instance, plot_iccs_seismograms + from aimbat.core import ( + create_iccs_instance, + plot_iccs_seismograms, + get_active_event, + ) from sqlmodel import Session with Session(engine) as session: - iccs = create_iccs_instance(session) + active_event = get_active_event(session) + iccs = create_iccs_instance(session, active_event).iccs plot_iccs_seismograms( iccs, iccs_parameters.context, iccs_parameters.all, return_fig=False ) diff --git a/src/aimbat/_cli/project.py b/src/aimbat/_cli/project.py index 3593745c..62dd311d 100644 --- a/src/aimbat/_cli/project.py +++ b/src/aimbat/_cli/project.py @@ -51,9 +51,81 @@ def cli_project_info( """Show information on an existing project.""" from aimbat.db import engine - from aimbat.core import print_project_info - - print_project_info(engine) + from aimbat.core._project import _project_exists + from aimbat.core import get_active_event + from aimbat.models import AimbatEvent, AimbatSeismogram, AimbatStation + from aimbat.logger import logger + from sqlmodel import Session, select + from rich.console import Console + from rich.table import Table + from rich.panel import Panel + from sqlalchemy.exc import NoResultFound + import aimbat.core._event as event + import aimbat.core._seismogram as seismogram + import aimbat.core._station as station + + logger.info("Printing project info.") + + if not _project_exists(engine): + raise RuntimeError( + 'No AIMBAT project found. Try running "aimbat project create" first.' + ) + + with Session(engine) as session: + grid = Table.grid(expand=False) + grid.add_column() + grid.add_column(justify="left") + if engine.driver == "pysqlite": + if engine.url.database == ":memory:": + grid.add_row("AIMBAT Project: ", "in-memory database") + else: + grid.add_row("AIMBAT Project File: ", str(engine.url.database)) + + events = len(session.exec(select(AimbatEvent)).all()) + completed_events = len(event.get_completed_events(session)) + stations = len(session.exec(select(AimbatStation)).all()) + seismograms = len(session.exec(select(AimbatSeismogram)).all()) + selected_seismograms = len( + seismogram.get_selected_seismograms(session, all_events=True) + ) + + grid.add_row( + "Number of Events (total/completed): ", + f"({events}/{completed_events})", + ) + + try: + active_event = get_active_event(session) + active_event_id = active_event.id + active_stations = len(station.get_stations_in_event(session, active_event)) + seismograms_in_event = len(active_event.seismograms) + selected_seismograms_in_event = len( + seismogram.get_selected_seismograms(session, event=active_event) + ) + except NoResultFound: + active_event_id = None + active_stations = None + seismograms_in_event = None + selected_seismograms_in_event = None + grid.add_row("Active Event ID: ", f"{active_event_id}") + grid.add_row( + "Number of Stations in Project (total/active event): ", + f"({stations}/{active_stations})", + ) + + grid.add_row( + "Number of Seismograms in Project (total/selected): ", + f"({seismograms}/{selected_seismograms})", + ) + grid.add_row( + "Number of Seismograms in Active Event (total/selected): ", + f"({seismograms_in_event}/{selected_seismograms_in_event})", + ) + + console = Console() + console.print( + Panel(grid, title="Project Info", title_align="left", border_style="dim") + ) if __name__ == "__main__": diff --git a/src/aimbat/_cli/seismogram.py b/src/aimbat/_cli/seismogram.py index be007cf2..df0e8536 100644 --- a/src/aimbat/_cli/seismogram.py +++ b/src/aimbat/_cli/seismogram.py @@ -55,23 +55,6 @@ def cli_seismogram_dump( print_json(dump_seismogram_table_to_json(session)) -@app.command(name="list") -@simple_exception -def cli_seismogram_list( - *, - all_events: Annotated[bool, ALL_EVENTS_PARAMETER] = False, - table_parameters: TableParameters = TableParameters(), - global_parameters: GlobalParameters = GlobalParameters(), -) -> None: - """Print information on the seismograms in the active event.""" - from aimbat.db import engine - from aimbat.core import print_seismogram_table - from sqlmodel import Session - - with Session(engine) as session: - print_seismogram_table(session, table_parameters.short, all_events) - - @parameter.command(name="get") @simple_exception def cli_seismogram_parameter_get( @@ -140,13 +123,16 @@ def cli_seismogram_parameter_dump( ) -> None: """Dump seismogram parameter table to json.""" from aimbat.db import engine - from aimbat.core import dump_seismogram_parameter_table_to_json + from aimbat.core import dump_seismogram_parameter_table_to_json, get_active_event from sqlmodel import Session from rich import print_json with Session(engine) as session: + active_event = get_active_event(session) if not all_events else None print_json( - dump_seismogram_parameter_table_to_json(session, all_events, as_string=True) + dump_seismogram_parameter_table_to_json( + session, all_events, as_string=True, event=active_event + ) ) @@ -163,11 +149,135 @@ def cli_seismogram_parameter_list( """ from aimbat.db import engine - from aimbat.core import print_seismogram_parameter_table + from aimbat.core import get_active_event, dump_seismogram_parameter_table_to_json + from aimbat.utils import uuid_shortener, json_to_table, TABLE_STYLING + from aimbat.logger import logger from sqlmodel import Session + short = table_parameters.short + with Session(engine) as session: - print_seismogram_parameter_table(session, table_parameters.short) + logger.info("Printing AIMBAT seismogram parameters table for active event.") + + active_event = get_active_event(session) + title = f"Seismogram parameters for event: {uuid_shortener(session, active_event) if short else str(active_event.id)}" + + json_to_table( + data=dump_seismogram_parameter_table_to_json( + session, all_events=False, as_string=False, event=active_event + ), + title=title, + skip_keys=["id"], + column_order=["seismogram_id", "select"], + common_column_kwargs={"highlight": True}, + formatters={ + "seismogram_id": lambda x: ( + uuid_shortener(session, AimbatSeismogram, str_uuid=x) + if short + else x + ), + }, + column_kwargs={ + "seismogram_id": { + "header": "Seismogram ID (shortened)" if short else "Seismogram ID", + "justify": "center", + "style": TABLE_STYLING.mine, + }, + }, + ) + + +@app.command(name="list") +@simple_exception +def cli_seismogram_list( + *, + all_events: Annotated[bool, ALL_EVENTS_PARAMETER] = False, + table_parameters: TableParameters = TableParameters(), + global_parameters: GlobalParameters = GlobalParameters(), +) -> None: + """Print information on the seismograms in the active event.""" + from aimbat.db import engine + from aimbat.core import get_active_event + from aimbat.utils import uuid_shortener, make_table, TABLE_STYLING + from aimbat.logger import logger + from rich.console import Console + from sqlmodel import Session, select + + short = table_parameters.short + + with Session(engine) as session: + logger.info("Printing AIMBAT seismogram table.") + + title = "AIMBAT seismograms for all events" + + if all_events: + logger.debug("Selecting seismograms for all events.") + seismograms = session.exec(select(AimbatSeismogram)).all() + else: + logger.debug("Selecting seismograms for active event only.") + active_event = get_active_event(session) + seismograms = active_event.seismograms + if short: + title = f"AIMBAT seismograms for event {active_event.time.strftime('%Y-%m-%d %H:%M:%S')} (ID={uuid_shortener(session, active_event)})" + else: + title = f"AIMBAT seismograms for event {active_event.time} (ID={active_event.id})" + + logger.debug(f"Found {len(seismograms)} seismograms for the table.") + + table = make_table(title=title) + table.add_column( + "ID (shortened)" if short else "ID", + justify="center", + style=TABLE_STYLING.id, + no_wrap=True, + ) + table.add_column( + "Selected", justify="center", style=TABLE_STYLING.mine, no_wrap=True + ) + table.add_column( + "NPTS", justify="center", style=TABLE_STYLING.mine, no_wrap=True + ) + table.add_column( + "Delta", justify="center", style=TABLE_STYLING.mine, no_wrap=True + ) + table.add_column( + "Data ID", justify="center", style=TABLE_STYLING.linked, no_wrap=True + ) + table.add_column("Station ID", justify="center", style=TABLE_STYLING.linked) + table.add_column("Station Name", justify="center", style=TABLE_STYLING.linked) + if all_events: + table.add_column("Event ID", justify="center", style=TABLE_STYLING.linked) + + for seismogram in seismograms: + logger.debug(f"Adding seismogram with ID {seismogram.id} to the table.") + row = [ + (uuid_shortener(session, seismogram) if short else str(seismogram.id)), + TABLE_STYLING.bool_formatter(seismogram.parameters.select), + str(len(seismogram.data)), + str(seismogram.delta.total_seconds()), + ( + uuid_shortener(session, seismogram.datasource) + if short + else str(seismogram.datasource.id) + ), + ( + uuid_shortener(session, seismogram.station) + if short + else str(seismogram.station.id) + ), + f"{seismogram.station.name} - {seismogram.station.network}", + ] + + if all_events: + row.append( + uuid_shortener(session, seismogram.event) + if short + else str(seismogram.event.id) + ) + table.add_row(*row) + + console = Console() + console.print(table) if __name__ == "__main__": diff --git a/src/aimbat/_cli/snapshot.py b/src/aimbat/_cli/snapshot.py index 10f39f02..a8c7c319 100644 --- a/src/aimbat/_cli/snapshot.py +++ b/src/aimbat/_cli/snapshot.py @@ -37,11 +37,12 @@ def cli_snapshot_create( comment: Optional description to help identify this snapshot later. """ from aimbat.db import engine - from aimbat.core import create_snapshot + from aimbat.core import create_snapshot, get_active_event from sqlmodel import Session with Session(engine) as session: - create_snapshot(session, comment) + active_event = get_active_event(session) + create_snapshot(session, active_event, comment) @app.command(name="rollback") @@ -84,12 +85,17 @@ def cli_snapshot_dump( ) -> None: """Dump the contents of the AIMBAT snapshot table to json.""" from aimbat.db import engine - from aimbat.core import dump_snapshot_tables_to_json + from aimbat.core import dump_snapshot_tables_to_json, get_active_event from sqlmodel import Session from rich import print_json with Session(engine) as session: - print_json(dump_snapshot_tables_to_json(session, all_events, as_string=True)) + active_event = get_active_event(session) if not all_events else None + print_json( + dump_snapshot_tables_to_json( + session, all_events, as_string=True, event=active_event + ) + ) @app.command(name="list") @@ -101,11 +107,86 @@ def cli_snapshot_list( ) -> None: """Print information on the snapshots for the active event.""" from aimbat.db import engine - from aimbat.core import print_snapshot_table + from aimbat.core import get_active_event, dump_snapshot_tables_to_json + from aimbat.utils import uuid_shortener, json_to_table, TABLE_STYLING + from aimbat.logger import logger + from aimbat.models import AimbatEvent + from pandas import Timestamp from sqlmodel import Session + short = table_parameters.short + with Session(engine) as session: - print_snapshot_table(session, table_parameters.short, all_events) + logger.info("Printing AIMBAT snapshots table.") + + title = "AIMBAT snapshots for all events" + + active_event = None + if not all_events: + active_event = get_active_event(session) + if short: + title = f"AIMBAT snapshots for event {active_event.time.strftime('%Y-%m-%d %H:%M:%S')} (ID={uuid_shortener(session, active_event)})" + else: + title = f"AIMBAT snapshots for event {active_event.time} (ID={active_event.id})" + + data = dump_snapshot_tables_to_json( + session, all_events, as_string=False, event=active_event + ) + snapshot_data = data["snapshots"] + + column_order = ["id", "date", "comment", "seismogram_count"] + if all_events: + column_order.append("event_id") + + skip_keys = [] if all_events else ["event_id"] + + json_to_table( + data=snapshot_data, + title=title, + column_order=column_order, + skip_keys=skip_keys, + formatters={ + "id": lambda x: ( + uuid_shortener(session, AimbatSnapshot, str_uuid=x) if short else x + ), + "date": lambda x: TABLE_STYLING.timestamp_formatter( + Timestamp(x), short + ), + "event_id": lambda x: ( + uuid_shortener(session, AimbatEvent, str_uuid=x) if short else x + ), + }, + common_column_kwargs={"justify": "center"}, + column_kwargs={ + "id": { + "header": "ID (shortened)" if short else "ID", + "style": TABLE_STYLING.id, + "no_wrap": True, + }, + "date": { + "header": "Date & Time", + "style": TABLE_STYLING.mine, + "no_wrap": True, + }, + "comment": {"style": TABLE_STYLING.mine}, + "seismogram_count": { + "header": "# Seismograms", + "style": TABLE_STYLING.linked, + }, + "selected_seismogram_count": { + "header": "# Selected", + "style": TABLE_STYLING.linked, + }, + "flipped_seismogram_count": { + "header": "# Flipped", + "style": TABLE_STYLING.linked, + }, + "event_id": { + "header": "Event ID (shortened)" if short else "Event ID", + "style": TABLE_STYLING.linked, + }, + }, + ) @app.command(name="details") @@ -117,12 +198,32 @@ def cli_snapshot_details( ) -> None: """Print information on the event parameters saved in a snapshot.""" from aimbat.db import engine - from aimbat.core import print_snapshot_parameters_table_by_id + from aimbat.utils import uuid_shortener, json_to_table, TABLE_STYLING from sqlmodel import Session + short = table_parameters.short + with Session(engine) as session: - print_snapshot_parameters_table_by_id( - session, snapshot_id, table_parameters.short + snapshot = session.get(AimbatSnapshot, snapshot_id) + + if snapshot is None: + raise ValueError( + f"Unable to print snapshot parameters: snapshot with id={snapshot_id} not found." + ) + + parameters_snapshot = snapshot.event_parameters_snapshot + json_to_table( + data=parameters_snapshot.model_dump(mode="json"), + title=f"Saved event parameters in snapshot: {uuid_shortener(session, parameters_snapshot.snapshot) if short else str(parameters_snapshot.snapshot.id)}", + skip_keys=["id", "snapshot_id", "parameters_id"], + common_column_kwargs={"highlight": True}, + column_kwargs={ + "Key": { + "header": "Parameter", + "justify": "left", + "style": TABLE_STYLING.id, + }, + }, ) diff --git a/src/aimbat/_cli/station.py b/src/aimbat/_cli/station.py index 842ad402..f4bb6af5 100644 --- a/src/aimbat/_cli/station.py +++ b/src/aimbat/_cli/station.py @@ -31,23 +31,6 @@ def cli_station_delete( delete_station_by_id(session, station_id) -@app.command(name="list") -@simple_exception -def cli_station_list( - *, - all_events: Annotated[bool, ALL_EVENTS_PARAMETER] = False, - table_parameters: TableParameters = TableParameters(), - global_parameters: GlobalParameters = GlobalParameters(), -) -> None: - """Print information on the stations used in the active event.""" - from aimbat.db import engine - from aimbat.core import print_station_table - from sqlmodel import Session - - with Session(engine) as session: - print_station_table(session, table_parameters.short, all_events) - - @app.command(name="dump") @simple_exception def cli_station_dump( @@ -68,5 +51,133 @@ def cli_station_dump( print_json(dump_station_table_to_json(session)) +@app.command(name="list") +@simple_exception +def cli_station_list( + *, + all_events: Annotated[bool, ALL_EVENTS_PARAMETER] = False, + table_parameters: TableParameters = TableParameters(), + global_parameters: GlobalParameters = GlobalParameters(), +) -> None: + """Print information on the stations used in the active event.""" + from aimbat.db import engine + from aimbat.core import ( + get_active_event, + get_stations_in_event, + ) + from aimbat.utils import uuid_shortener, json_to_table, TABLE_STYLING + from aimbat.logger import logger + from typing import Any + from sqlmodel import Session + + short = table_parameters.short + + with Session(engine) as session: + logger.info("Printing station table.") + + title = "AIMBAT stations for all events" + + if all_events: + logger.debug("Selecting all AIMBAT stations.") + from aimbat.core import dump_station_table_with_counts + + data = dump_station_table_with_counts(session) + else: + logger.debug("Selecting AIMBAT stations used by active event.") + active_event = get_active_event(session) + data = get_stations_in_event(session, active_event, as_json=True) + + if short: + title = f"AIMBAT stations for event {active_event.time.strftime('%Y-%m-%d %H:%M:%S')} (ID={uuid_shortener(session, active_event)})" + else: + title = f"AIMBAT stations for event {active_event.time} (ID={active_event.id})" + + column_order = [ + "id", + "name", + "network", + "channel", + "location", + "latitude", + "longitude", + "elevation", + ] + if all_events: + column_order.extend(["seismogram_count", "event_count"]) + + column_kwargs: dict[str, dict[str, Any]] = { + "id": { + "header": "ID (shortened)" if short else "ID", + "style": TABLE_STYLING.id, + "justify": "center", + "no_wrap": True, + }, + "name": { + "header": "Name", + "style": TABLE_STYLING.mine, + "justify": "center", + "no_wrap": True, + }, + "network": { + "header": "Network", + "style": TABLE_STYLING.mine, + "justify": "center", + "no_wrap": True, + }, + "channel": { + "header": "Channel", + "style": TABLE_STYLING.mine, + "justify": "center", + }, + "location": { + "header": "Location", + "style": TABLE_STYLING.mine, + "justify": "center", + }, + "latitude": { + "header": "Latitude", + "style": TABLE_STYLING.mine, + "justify": "center", + }, + "longitude": { + "header": "Longitude", + "style": TABLE_STYLING.mine, + "justify": "center", + }, + "elevation": { + "header": "Elevation", + "style": TABLE_STYLING.mine, + "justify": "center", + }, + "seismogram_count": { + "header": "# Seismograms", + "style": TABLE_STYLING.linked, + "justify": "center", + }, + "event_count": { + "header": "# Events", + "style": TABLE_STYLING.linked, + "justify": "center", + }, + } + + formatters = { + "id": lambda x: ( + uuid_shortener(session, AimbatStation, str_uuid=x) if short else str(x) + ), + "latitude": lambda x: f"{x:.3f}" if short else str(x), + "longitude": lambda x: f"{x:.3f}" if short else str(x), + "elevation": lambda x: f"{x:.0f}" if short else str(x), + } + + json_to_table( + data, + title=title, + column_order=column_order, + column_kwargs=column_kwargs, + formatters=formatters, + ) + + if __name__ == "__main__": app() diff --git a/src/aimbat/_cli/utils/sampledata.py b/src/aimbat/_cli/utils/sampledata.py index 34c7ac18..d3fc352b 100644 --- a/src/aimbat/_cli/utils/sampledata.py +++ b/src/aimbat/_cli/utils/sampledata.py @@ -17,7 +17,7 @@ @app.command(name="download") @simple_exception def sampledata_cli_download( - *, force: bool = False, global_parameters: GlobalParameters | None = None + *, force: bool = False, global_parameters: GlobalParameters = GlobalParameters() ) -> None: """Download AIMBAT sample data. @@ -29,19 +29,17 @@ def sampledata_cli_download( """ from aimbat.utils import download_sampledata - global_parameters = global_parameters or GlobalParameters() - download_sampledata(force) @app.command(name="delete") @simple_exception -def sampledata_cli_delete(*, global_parameters: GlobalParameters | None = None) -> None: +def sampledata_cli_delete( + *, global_parameters: GlobalParameters = GlobalParameters() +) -> None: """Recursively delete sample data directory.""" from aimbat.utils import delete_sampledata - global_parameters = global_parameters or GlobalParameters() - delete_sampledata() diff --git a/src/aimbat/_tui/app.py b/src/aimbat/_tui/app.py index c9a3fa7d..0ec0cd73 100644 --- a/src/aimbat/_tui/app.py +++ b/src/aimbat/_tui/app.py @@ -10,7 +10,7 @@ from rich.console import Console from rich.panel import Panel -from pandas import Timedelta +from pandas import Timedelta, Timestamp from pysmo.tools.iccs import ICCS from sqlalchemy.exc import NoResultFound from sqlmodel import Session @@ -32,6 +32,7 @@ from aimbat._types import EventParameter, SeismogramParameter from aimbat.models._parameters import AimbatEventParametersBase from aimbat.core import ( + BoundICCS, add_data_to_project, create_iccs_instance, create_snapshot, @@ -125,7 +126,8 @@ def compose(self) -> ComposeResult: yield Footer() def on_mount(self) -> None: - self._iccs: ICCS | None = None + self._bound_iccs: BoundICCS | None = None + self._iccs_last_modified_seen: Timestamp | None = None self._active_tab: str = "tab-seismograms" self.theme = _DEFAULT_THEME @@ -135,6 +137,7 @@ def on_mount(self) -> None: self._setup_station_table() self._setup_snapshot_table() + self.set_interval(5, self._check_iccs_staleness) self._create_iccs() self.refresh_all() @@ -164,7 +167,7 @@ def _create_iccs(self) -> None: ICCS construction reads waveform data, so it must not block the asyncio event loop. """ - self._iccs = None + self._bound_iccs = None self._worker_create_iccs() @work(thread=True) @@ -172,7 +175,8 @@ def _worker_create_iccs(self) -> None: """Background worker: create ICCS instance without blocking the UI.""" try: with Session(engine) as session: - new_iccs = create_iccs_instance(session) + active_event = get_active_event(session) + bound_iccs = create_iccs_instance(session, active_event) except (NoResultFound, RuntimeError): return except Exception as exc: @@ -180,11 +184,11 @@ def _worker_create_iccs(self) -> None: self.notify, f"ICCS init failed: {exc}", severity="error" ) return - self.call_from_thread(self._assign_iccs, new_iccs) + self.call_from_thread(self._assign_iccs, bound_iccs) - def _assign_iccs(self, iccs: ICCS) -> None: - """Main-thread callback: store the new ICCS instance and refresh status.""" - self._iccs = iccs + def _assign_iccs(self, bound_iccs: BoundICCS) -> None: + """Main-thread callback: store the new BoundICCS instance and refresh status.""" + self._bound_iccs = bound_iccs self._refresh_event_bar() self._refresh_seismograms() @@ -227,22 +231,49 @@ def refresh_all(self) -> None: self._refresh_stations() self._refresh_snapshots() + def _check_iccs_staleness(self) -> None: + """Trigger ICCS recreation if the active event has been modified externally. + + When ICCS creation previously failed (e.g. due to an invalid parameter set via + the CLI), retries whenever ``event.last_modified`` changes. On any detected + change the full UI is refreshed so panels reflect the new DB state immediately. + """ + try: + with Session(engine) as session: + event = get_active_event(session) + changed = False + if self._bound_iccs is not None: + if self._bound_iccs.is_stale(event): + self._iccs_last_modified_seen = event.last_modified + self._create_iccs() + changed = True + elif event.last_modified != self._iccs_last_modified_seen: + self._iccs_last_modified_seen = event.last_modified + self._create_iccs() + changed = True + except (NoResultFound, RuntimeError): + return + if changed: + self.refresh_all() + def _refresh_event_bar(self) -> None: bar = self.query_one("#event-bar", Static) - iccs_status = " ● ICCS ready" if self._iccs is not None else " ○ no ICCS" try: with Session(engine) as session: event = get_active_event(session) + iccs_status = ( + " ● ICCS ready" if self._bound_iccs is not None else " ○ no ICCS" + ) time_str = str(event.time)[:19] if event.time else "unknown" lat = f"{event.latitude:.3f}°" if event.latitude is not None else "?" lon = f"{event.longitude:.3f}°" if event.longitude is not None else "?" - depth = ( - f" depth {event.depth / 1000:.1f} km" - if event.depth is not None + modified = ( + f" modified: {str(event.last_modified)[:19]}" + if event.last_modified is not None else "" ) bar.update( - f"Active event: {time_str} | {lat}, {lon}{depth}" + f"Active event: {time_str} | {lat}, {lon}{modified}" f" [dim]{iccs_status} e = switch event[/dim]" ) except NoResultFound: @@ -256,10 +287,10 @@ def _refresh_seismograms(self) -> None: table.clear() ccnorm_map: dict[uuid.UUID, float] = {} - if self._iccs is not None: + if self._bound_iccs is not None: try: for iccs_seis, ccnorm in zip( - self._iccs.seismograms, self._iccs.ccnorms + self._bound_iccs.iccs.seismograms, self._bound_iccs.iccs.ccnorms ): ccnorm_map[iccs_seis.extra["id"]] = float(ccnorm) except Exception: @@ -467,23 +498,27 @@ def on_input(raw: str | None) -> None: def _apply_parameter(self, attr: str, value: object) -> None: """Write a parameter to the DB and sync to the in-memory ICCS object.""" + iccs = self._bound_iccs.iccs if self._bound_iccs is not None else None + # Validate with ICCS first — before touching the DB — so invalid values # are rejected without being persisted. - if self._iccs is not None and hasattr(self._iccs, attr): + if iccs is not None and hasattr(iccs, attr): try: - setattr(self._iccs, attr, value) - self._iccs.clear_cache() + setattr(iccs, attr, value) + iccs.clear_cache() except ValueError as exc: self.notify(str(exc), severity="error") return try: with Session(engine) as session: + active_event = get_active_event(session) if attr in {p.value for p in EventParameter}: - set_event_parameter(session, EventParameter(attr), value) # type: ignore[call-overload] + set_event_parameter( + session, active_event, EventParameter(attr), value + ) # type: ignore[call-overload] else: # mccc_damp / mccc_min_ccnorm — not in EventParameter enum - active_event = get_active_event(session) validated = AimbatEventParametersBase.model_validate( active_event.parameters, update={attr: value} ) @@ -502,9 +537,12 @@ def _apply_parameter(self, attr: str, value: object) -> None: self._create_iccs() # revert ICCS to DB state return - # Parameter change may have fixed previously invalid ranges. - if self._iccs is None: + if self._bound_iccs is None: + # Parameter change may have fixed previously invalid ranges. self._create_iccs() + else: + # Acknowledge our own write so staleness check doesn't recreate. + self._bound_iccs.created_at = Timestamp.now("UTC") self._refresh_parameters() self._refresh_seismograms() @@ -550,11 +588,12 @@ def _toggle_seismogram_bool(self, item_id: str, param: SeismogramParameter) -> N setattr(seis.parameters, param, new_value) session.add(seis) session.commit() - if self._iccs is not None: - for iccs_seis in self._iccs.seismograms: + if self._bound_iccs is not None: + for iccs_seis in self._bound_iccs.iccs.seismograms: if iccs_seis.extra.get("id") == seis_uuid: setattr(iccs_seis, param, new_value) - self._iccs.clear_cache() + self._bound_iccs.iccs.clear_cache() + self._bound_iccs.created_at = Timestamp.now("UTC") break self._refresh_seismograms() self.notify(f"{param} toggled", timeout=2) @@ -635,9 +674,13 @@ def on_confirm(confirmed: bool | None) -> None: try: with Session(engine) as session: rollback_to_snapshot_by_id(session, uuid.UUID(snap_id)) - if self._iccs is not None: - sync_iccs_parameters(session, self._iccs) - if self._iccs is None: + if self._bound_iccs is not None: + active_event = get_active_event(session) + sync_iccs_parameters( + session, active_event, self._bound_iccs.iccs + ) + self._bound_iccs.created_at = Timestamp.now("UTC") + if self._bound_iccs is None: self._create_iccs() self.refresh_all() self.notify("Rolled back to snapshot", timeout=3) @@ -696,9 +739,23 @@ def on_file(path: Path | None) -> None: self.push_screen(ActionMenuModal("Add Data", actions), on_type) + def _require_iccs(self) -> bool: + """Return True if ICCS is ready; show a contextual warning and return False otherwise.""" + if self._bound_iccs is not None: + return True + try: + with Session(engine) as session: + get_active_event(session) + self.notify( + "ICCS not ready — check event parameters (Parameters tab)", + severity="warning", + ) + except (NoResultFound, RuntimeError): + self.notify("No active event — press e to select one", severity="warning") + return False + def action_open_interactive_tools(self) -> None: - if self._iccs is None: - self.notify("No active event — activate one first", severity="warning") + if not self._require_iccs(): return def on_result(result: tuple[str, bool, bool] | None) -> None: @@ -714,7 +771,7 @@ def _run_pick_tool(self, tool: str, context: bool, all_seis: bool) -> None: matplotlib on the main thread via App.suspend(), which is the correct Textual pattern for blocking terminal-adjacent processes. """ - if self._iccs is None: + if self._bound_iccs is None: self.notify("ICCS not ready — please wait", severity="warning") return _TOOL_LABELS = { @@ -738,10 +795,11 @@ def _run_pick_tool(self, tool: str, context: bool, all_seis: bool) -> None: ) ) with Session(engine) as session: + active_event = get_active_event(session) if tool == "phase": update_pick( session, - self._iccs, + self._bound_iccs.iccs, context, all_seis, False, @@ -750,7 +808,8 @@ def _run_pick_tool(self, tool: str, context: bool, all_seis: bool) -> None: elif tool == "window": update_timewindow( session, - self._iccs, + active_event, + self._bound_iccs.iccs, context, all_seis, False, @@ -759,7 +818,8 @@ def _run_pick_tool(self, tool: str, context: bool, all_seis: bool) -> None: elif tool == "ccnorm": update_min_ccnorm( session, - self._iccs, + active_event, + self._bound_iccs.iccs, context, all_seis, return_fig=False, @@ -768,19 +828,19 @@ def _run_pick_tool(self, tool: str, context: bool, all_seis: bool) -> None: except Exception as exc: self.notify(str(exc), severity="error") return + self._bound_iccs.created_at = Timestamp.now("UTC") self._refresh_parameters() self._refresh_seismograms() self._refresh_event_bar() self.notify("Done", timeout=2) def action_open_align(self) -> None: - if self._iccs is None: - self.notify("No active event — activate one first", severity="warning") + if not self._require_iccs(): return def on_result(result: tuple[str, bool, bool, bool] | None) -> None: if result is not None: - self._run_align_tool(self._iccs, *result) + self._run_align_tool(self._bound_iccs.iccs, *result) # type: ignore[union-attr] self.push_screen(AlignModal(), on_result) @@ -799,13 +859,18 @@ def _run_align_tool( if algorithm == "iccs": run_iccs(session, iccs, autoflip, autoselect) elif algorithm == "mccc": - run_mccc(session, iccs, all_seis) + active_event = get_active_event(session) + run_mccc(session, active_event, iccs, all_seis) except Exception as exc: self.call_from_thread(self.notify, str(exc), severity="error") return self.call_from_thread(self._post_align_complete) def _post_align_complete(self) -> None: + # Acknowledge our own writes (t1/flip/select written back by ICCS/MCCC) + # so the staleness check doesn't recreate an ICCS we just ran. + if self._bound_iccs is not None: + self._bound_iccs.created_at = Timestamp.now("UTC") self.refresh_all() self.notify("Alignment complete", timeout=3) @@ -815,7 +880,8 @@ def on_comment(comment: str | None) -> None: return try: with Session(engine) as session: - create_snapshot(session, comment or None) + active_event = get_active_event(session) + create_snapshot(session, active_event, comment or None) self._refresh_snapshots() self.notify("Snapshot created", timeout=2) except Exception as exc: diff --git a/src/aimbat/_types/_event.py b/src/aimbat/_types/_event.py index 66aa3937..efd7998c 100644 --- a/src/aimbat/_types/_event.py +++ b/src/aimbat/_types/_event.py @@ -3,10 +3,10 @@ class EventParameter(StrEnum): - """[`AimbatEvent`][aimbat.lib.models.AimbatEvent] enum class for typing. + """[`AimbatEvent`][aimbat.models.AimbatEvent] enum class for typing. This enum class is used for typing, cli args etc. The attributes must be - the same as in the [`AimbatEvent`][aimbat.lib.models.AimbatEvent] model. + the same as in the [`AimbatEvent`][aimbat.models.AimbatEvent] model. """ COMPLETED = auto() @@ -21,16 +21,16 @@ class EventParameter(StrEnum): type EventParameterBool = Literal[ EventParameter.COMPLETED, EventParameter.BANDPASS_APPLY ] -"[`TypeAlias`][typing.TypeAlias] for [`AimbatEvent`][aimbat.lib.models.AimbatEvent] attributes with [`bool`][bool] values." +"[`TypeAlias`][typing.TypeAlias] for [`AimbatEvent`][aimbat.models.AimbatEvent] attributes with [`bool`][bool] values." type EventParameterFloat = Literal[ EventParameter.MIN_CCNORM, EventParameter.BANDPASS_FMIN, EventParameter.BANDPASS_FMAX, ] -"[`TypeAlias`][typing.TypeAlias] for [`AimbatEvent`][aimbat.lib.models.AimbatEvent] attributes with [`float`][float] values." +"[`TypeAlias`][typing.TypeAlias] for [`AimbatEvent`][aimbat.models.AimbatEvent] attributes with [`float`][float] values." type EventParameterTimedelta = Literal[ EventParameter.WINDOW_PRE, EventParameter.WINDOW_POST ] -"[`TypeAlias`][typing.TypeAlias] for [`AimbatEvent`][aimbat.lib.models.AimbatEvent] attributes with [`Timedelta`][pandas.Timedelta] values." +"[`TypeAlias`][typing.TypeAlias] for [`AimbatEvent`][aimbat.models.AimbatEvent] attributes with [`Timedelta`][pandas.Timedelta] values." diff --git a/src/aimbat/_types/_seismogram.py b/src/aimbat/_types/_seismogram.py index f59bd64f..d8561f2f 100644 --- a/src/aimbat/_types/_seismogram.py +++ b/src/aimbat/_types/_seismogram.py @@ -9,10 +9,10 @@ class SeismogramParameter(StrEnum): - """[`AimbatSeismograParameters`][aimbat.lib.models.AimbatSeismogramParameters] enum class for typing. + """[`AimbatSeismograParameters`][aimbat.models.AimbatSeismogramParameters] enum class for typing. This enum class is used for typing, cli args etc. The attributes must be - the same as in the [`AimbatParameters`][aimbat.lib.models.AimbatParameters] model. + the same as in the [`AimbatParameters`][aimbat.models.AimbatParameters] model. """ SELECT = auto() diff --git a/src/aimbat/core/_active_event.py b/src/aimbat/core/_active_event.py index ae318a10..87a21767 100644 --- a/src/aimbat/core/_active_event.py +++ b/src/aimbat/core/_active_event.py @@ -1,14 +1,14 @@ """Get and set the active event (i.e. the one being processed).""" # WARNING: Do not import other modules from `aimbat.core` here to avoid circular imports -from aimbat.io import clear_seismogram_cache -from aimbat.logger import logger -from aimbat.models import AimbatEvent -from aimbat._cli.common import HINTS from sqlmodel import Session, select from sqlalchemy.exc import NoResultFound from contextlib import suppress from uuid import UUID +from aimbat.io import clear_seismogram_cache +from aimbat.logger import logger +from aimbat.models import AimbatEvent +from aimbat._cli.common import HINTS __all__ = [ "get_active_event", diff --git a/src/aimbat/core/_data.py b/src/aimbat/core/_data.py index 147c49e3..3fccb38c 100644 --- a/src/aimbat/core/_data.py +++ b/src/aimbat/core/_data.py @@ -1,11 +1,13 @@ import os import uuid -from aimbat.core import get_active_event +from sqlmodel import Session, select +from pydantic import TypeAdapter +from collections.abc import Sequence +from rich.progress import track +from rich.console import Console from aimbat.logger import logger from aimbat.io import DataType from aimbat.utils import ( - uuid_shortener, - make_table, TABLE_STYLING, json_to_table, ) @@ -24,16 +26,10 @@ AimbatEvent, AimbatSeismogram, ) -from sqlmodel import Session, select -from pydantic import TypeAdapter -from collections.abc import Sequence -from rich.progress import track -from rich.console import Console __all__ = [ "add_data_to_project", - "get_data_for_active_event", - "print_data_table", + "get_data_for_event", "dump_data_table_to_json", ] @@ -315,84 +311,29 @@ def add_data_to_project( raise -def get_data_for_active_event(session: Session) -> Sequence[AimbatDataSource]: - """Returns the data sources belonging to the active event. +def get_data_for_event( + session: Session, event: AimbatEvent +) -> Sequence[AimbatDataSource]: + """Returns the data sources belonging to the given event. Args: session: Database session. + event: AimbatEvent. Returns: - Sequence of AimbatDataSource objects belonging to the active event. + Sequence of AimbatDataSource objects belonging to the event. """ - logger.info("Getting data sources for active event.") + logger.info(f"Getting data sources for event {event.id}.") statement = ( select(AimbatDataSource) .join(AimbatSeismogram) - .join(AimbatEvent) - .where(AimbatEvent.active == 1) + .where(AimbatSeismogram.event_id == event.id) ) return session.exec(statement).all() -def print_data_table(session: Session, short: bool, all_events: bool = False) -> None: - """Print a pretty table with information about the data sources in the database. - - Args: - short: Shorten UUIDs and format data. - all_events: Print all data sources instead of limiting to the active event. - """ - - logger.info("Printing data sources table.") - - if all_events: - aimbat_data_sources = session.exec(select(AimbatDataSource)).all() - title = "Data sources for all events" - else: - active_event = get_active_event(session) - aimbat_data_sources = get_data_for_active_event(session) - time = ( - active_event.time.strftime("%Y-%m-%d %H:%M:%S") - if short - else active_event.time - ) - id = uuid_shortener(session, active_event) if short else active_event.id - title = f"Data sources for event {time} (ID={id})" - - logger.debug(f"Found {len(aimbat_data_sources)} data sources in total.") - - rows = [ - [ - uuid_shortener(session, a) if short else str(a.id), - str(a.datatype), - str(a.sourcename), - (uuid_shortener(session, a.seismogram) if short else str(a.seismogram.id)), - ] - for a in aimbat_data_sources - ] - - table = make_table(title=title) - - table.add_column( - "ID (shortened)" if short else "ID", - justify="center", - style=TABLE_STYLING.id, - no_wrap=True, - ) - table.add_column("Datatype", justify="center", style=TABLE_STYLING.mine) - table.add_column("Source", justify="left", style=TABLE_STYLING.mine, no_wrap=True) - table.add_column( - "Seismogram ID", justify="center", style=TABLE_STYLING.linked, no_wrap=True - ) - - for row in rows: - table.add_row(*row) - - console = Console() - console.print(table) - - def dump_data_table_to_json(session: Session) -> str: """Dump the table data to json.""" diff --git a/src/aimbat/core/_event.py b/src/aimbat/core/_event.py index 23a615f7..c5164359 100644 --- a/src/aimbat/core/_event.py +++ b/src/aimbat/core/_event.py @@ -1,13 +1,14 @@ """Module to manage and view events in AIMBAT.""" -from aimbat.core import get_active_event +from pydantic import TypeAdapter +from sqlmodel import select, Session +from sqlalchemy.exc import NoResultFound +from typing import overload, Any, Literal +from pandas import Timedelta +from collections.abc import Sequence +from uuid import UUID from aimbat.logger import logger from aimbat._cli.common import HINTS -from aimbat.utils import ( - uuid_shortener, - json_to_table, - TABLE_STYLING, -) from aimbat.models import ( AimbatEvent, AimbatEventParameters, @@ -22,13 +23,6 @@ EventParameterFloat, EventParameterTimedelta, ) -from pydantic import TypeAdapter -from sqlmodel import select, Session -from sqlalchemy.exc import NoResultFound -from typing import overload, Any, Literal -from pandas import Timedelta, Timestamp -from collections.abc import Sequence -from uuid import UUID __all__ = [ "delete_event_by_id", @@ -38,9 +32,7 @@ "get_event_parameter", "set_event_parameter", "dump_event_table_to_json", - "print_event_table", "dump_event_parameter_table_to_json", - "print_event_parameter_table", ] @@ -125,85 +117,96 @@ def get_events_using_station( @overload def get_event_parameter( - session: Session, name: EventParameterTimedelta + session: Session, event: AimbatEvent, name: EventParameterTimedelta ) -> Timedelta: ... @overload -def get_event_parameter(session: Session, name: EventParameterBool) -> bool: ... +def get_event_parameter( + session: Session, event: AimbatEvent, name: EventParameterBool +) -> bool: ... @overload -def get_event_parameter(session: Session, name: EventParameterFloat) -> float: ... +def get_event_parameter( + session: Session, event: AimbatEvent, name: EventParameterFloat +) -> float: ... @overload def get_event_parameter( - session: Session, name: EventParameter + session: Session, event: AimbatEvent, name: EventParameter ) -> Timedelta | bool | float: ... def get_event_parameter( - session: Session, name: EventParameter + session: Session, event: AimbatEvent, name: EventParameter ) -> Timedelta | bool | float: - """Get event parameter value for the active event. + """Get event parameter value for the given event. Args: session: Database session. + event: AimbatEvent. name: Name of the parameter. """ - active_event = get_active_event(session) + logger.info(f"Getting {name=} value for {event=}.") - logger.info(f"Getting {name=} value for {active_event=}.") - - return getattr(active_event.parameters, name) + return getattr(event.parameters, name) @overload def set_event_parameter( - session: Session, name: EventParameterTimedelta, value: Timedelta + session: Session, + event: AimbatEvent, + name: EventParameterTimedelta, + value: Timedelta, ) -> None: ... @overload def set_event_parameter( - session: Session, name: EventParameterFloat, value: float + session: Session, event: AimbatEvent, name: EventParameterFloat, value: float ) -> None: ... @overload def set_event_parameter( - session: Session, name: EventParameterBool, value: bool | str + session: Session, event: AimbatEvent, name: EventParameterBool, value: bool | str ) -> None: ... @overload def set_event_parameter( - session: Session, name: EventParameter, value: Timedelta | bool | float | str + session: Session, + event: AimbatEvent, + name: EventParameter, + value: Timedelta | bool | float | str, ) -> None: ... def set_event_parameter( - session: Session, name: EventParameter, value: Timedelta | bool | float | str + session: Session, + event: AimbatEvent, + name: EventParameter, + value: Timedelta | bool | float | str, ) -> None: - """Set event parameter value for the active event. + """Set event parameter value for the given event. Args: session: Database session. + event: AimbatEvent. name: Name of the parameter. value: Value to set. """ - active_event = get_active_event(session) - - logger.info(f"Setting {name=} to {value} for {active_event=}.") + logger.info(f"Setting {name=} to {value} for {event=}.") parameters = AimbatEventParametersBase.model_validate( - active_event.parameters, update={name: value} + event.parameters, update={name: value} ) - setattr(active_event.parameters, name, getattr(parameters, name)) - session.add(active_event) + setattr(event.parameters, name, getattr(parameters, name)) + session.add(event) session.commit() @@ -237,26 +240,45 @@ def dump_event_table_to_json( @overload def dump_event_parameter_table_to_json( - session: Session, all_events: bool, as_string: Literal[True] + session: Session, + all_events: bool, + as_string: Literal[True], + event: AimbatEvent | None = None, ) -> str: ... @overload def dump_event_parameter_table_to_json( - session: Session, all_events: Literal[False], as_string: Literal[False] + session: Session, + all_events: Literal[False], + as_string: Literal[False], + event: AimbatEvent | None = None, ) -> dict[str, Any]: ... @overload def dump_event_parameter_table_to_json( - session: Session, all_events: Literal[True], as_string: Literal[False] + session: Session, + all_events: Literal[True], + as_string: Literal[False], + event: AimbatEvent | None = None, ) -> list[dict[str, Any]]: ... def dump_event_parameter_table_to_json( - session: Session, all_events: bool, as_string: bool + session: Session, + all_events: bool, + as_string: bool, + event: AimbatEvent | None = None, ) -> str | dict[str, Any] | list[dict[str, Any]]: - """Dump the event parameter table data to json.""" + """Dump the event parameter table data to json. + + Args: + session: Database session. + all_events: Include event parameter table data for all events. + as_string: Whether to return the result as a string. + event: Event to dump parameter data for (only used when all_events is False). + """ logger.info("Dumping AIMBAT event parameter table to json.") @@ -270,130 +292,9 @@ def dump_event_parameter_table_to_json( else: return adapter.dump_python(parameters, mode="json") - active_event = get_active_event(session) + if event is None: + raise ValueError("An event must be provided when all_events is False.") if as_string: - return active_event.parameters.model_dump_json() - return active_event.parameters.model_dump(mode="json") - - -def print_event_table(session: Session, short: bool) -> None: - """Print a pretty table with AIMBAT events. - - Args: - session: Database session. - short: Shorten and format the output to be more human-readable. - """ - - logger.info("Printing AIMBAT events table.") - - json_to_table( - data=dump_event_table_to_json(session, as_string=False), - title="AIMBAT Events", - column_order=[ - "id", - "active", - "time", - "latitude", - "longitude", - "depth", - "completed", - "seismogram_count", - "station_count", - ], - formatters={ - "id": lambda x: ( - uuid_shortener(session, AimbatEvent, str_uuid=x) if short else x - ), - "active": TABLE_STYLING.bool_formatter, - "time": lambda x: TABLE_STYLING.timestamp_formatter(Timestamp(x), short), - "latitude": lambda x: f"{x:.3f}" if short else str(x), - "longitude": lambda x: f"{x:.3f}" if short else str(x), - "depth": lambda x: f"{x:.0f}" if short and x is not None else str(x), - "completed": TABLE_STYLING.bool_formatter, - }, - common_column_kwargs={"justify": "center"}, - column_kwargs={ - "id": { - "header": "ID (shortened)" if short else "ID", - "style": TABLE_STYLING.id, - "no_wrap": True, - }, - "active": {"style": TABLE_STYLING.mine, "no_wrap": True}, - "time": { - "header": "Date & Time", - "style": TABLE_STYLING.mine, - "no_wrap": True, - }, - "latitude": {"style": TABLE_STYLING.mine}, - "longitude": {"style": TABLE_STYLING.mine}, - "depth": {"style": TABLE_STYLING.mine}, - "completed": {"style": TABLE_STYLING.parameters}, - "seismogram_count": { - "header": "# Seismograms", - "style": TABLE_STYLING.linked, - }, - "station_count": { - "header": "# Stations", - "style": TABLE_STYLING.linked, - }, - }, - ) - - -def print_event_parameter_table( - session: Session, short: bool, all_events: bool -) -> None: - """Print a pretty table with AIMBAT parameter values for the active event. - - Args: - short: Shorten and format the output to be more human-readable. - all_events: Whether to print parameters for all events or just the active one. - """ - - if all_events: - logger.info("Printing AIMBAT event parameters table for all events.") - json_to_table( - data=dump_event_parameter_table_to_json( - session, all_events=True, as_string=False - ), - title="Event parameters for all events", - skip_keys=["id"], - column_order=[ - "event_id", - "completed", - "window_pre", - "window_post", - "min_ccnorm", - ], - formatters={ - "event_id": lambda x: ( - uuid_shortener(session, AimbatEvent, str_uuid=x) if short else x - ), - }, - common_column_kwargs={"highlight": True}, - column_kwargs={ - "event_id": { - "header": "Event ID (shortened)" if short else "Event ID", - "justify": "center", - "style": TABLE_STYLING.mine, - }, - }, - ) - else: - logger.info("Printing AIMBAT event parameters table for active event.") - - active_event = get_active_event(session) - json_to_table( - data=active_event.parameters.model_dump(mode="json"), - title=f"Event parameters for event: {uuid_shortener(session, active_event) if short else str(active_event.id)}", - skip_keys=["id", "event_id"], - common_column_kwargs={"highlight": True}, - column_kwargs={ - "Key": { - "header": "Parameter", - "justify": "left", - "style": TABLE_STYLING.id, - }, - }, - ) + return event.parameters.model_dump_json() + return event.parameters.model_dump(mode="json") diff --git a/src/aimbat/core/_iccs.py b/src/aimbat/core/_iccs.py index 8ad3bf0a..31c8ac7b 100644 --- a/src/aimbat/core/_iccs.py +++ b/src/aimbat/core/_iccs.py @@ -1,13 +1,11 @@ """Processing of data for AIMBAT.""" -from aimbat.core import get_active_event -from aimbat import settings -from aimbat.logger import logger -from aimbat.models import AimbatSeismogram -from aimbat.models._parameters import ( - AimbatEventParametersBase, - AimbatSeismogramParametersBase, -) +from dataclasses import dataclass +from uuid import UUID + +from pandas import Timestamp +from sqlmodel import Session +from pysmo.functions import clone_to_mini from pysmo.tools.iccs import ( ICCS, MiniICCSSeismogram, @@ -17,10 +15,16 @@ update_pick as _update_pick, update_timewindow as _update_timewindow, ) -from pysmo.functions import clone_to_mini -from sqlmodel import Session +from aimbat import settings +from aimbat.logger import logger +from aimbat.models import AimbatSeismogram, AimbatEvent +from aimbat.models._parameters import ( + AimbatEventParametersBase, + AimbatSeismogramParametersBase, +) __all__ = [ + "BoundICCS", "create_iccs_instance", "sync_iccs_parameters", "run_iccs", @@ -33,30 +37,55 @@ ] -def create_iccs_instance(session: Session) -> ICCS: - """Create an ICCS instance for the active event. +@dataclass +class BoundICCS: + """An ICCS instance explicitly bound to a specific event. + + Use `is_stale` to detect whether the event's parameters have been modified + (e.g. by a CLI command) since this instance was created. + """ + + iccs: ICCS + event_id: UUID + created_at: Timestamp + + def is_stale(self, event: AimbatEvent) -> bool: + """Return True if the event has been modified since this ICCS was created. + + Args: + event: The event to check against. + """ + if event.id != self.event_id: + return True + if event.last_modified is None: + return False + return event.last_modified > self.created_at + + +def create_iccs_instance(session: Session, event: AimbatEvent) -> BoundICCS: + """Create a BoundICCS instance for the given event. Seismogram data is copied into MiniICCSSeismogram objects so the session does not need to remain open after this call. Args: session: Database session. + event: AimbatEvent. Returns: - ICCS instance. + BoundICCS instance tied to the given event. """ - logger.info("Creating ICCS instance for active event.") + logger.info(f"Creating ICCS instance for event {event.id}.") - active_event = get_active_event(session) - p = active_event.parameters + p = event.parameters seismograms = [ clone_to_mini(MiniICCSSeismogram, seis, update={"extra": {"id": seis.id}}) - for seis in active_event.seismograms + for seis in event.seismograms ] - return ICCS( + iccs = ICCS( seismograms=seismograms, window_pre=p.window_pre, window_post=p.window_post, @@ -66,6 +95,7 @@ def create_iccs_instance(session: Session) -> ICCS: min_ccnorm=p.min_ccnorm, context_width=settings.context_width, ) + return BoundICCS(iccs=iccs, event_id=event.id, created_at=Timestamp.now("UTC")) def _write_back_seismograms(session: Session, iccs: ICCS) -> None: @@ -84,7 +114,7 @@ def _write_back_seismograms(session: Session, iccs: ICCS) -> None: session.commit() -def sync_iccs_parameters(session: Session, iccs: ICCS) -> None: +def sync_iccs_parameters(session: Session, event: AimbatEvent, iccs: ICCS) -> None: """Sync an existing ICCS instance's parameters from the database. Updates event-level and per-seismogram parameters without re-reading waveform @@ -93,13 +123,13 @@ def sync_iccs_parameters(session: Session, iccs: ICCS) -> None: Args: session: Database session. + event: AimbatEvent. iccs: ICCS instance to update in-place. """ - logger.info("Syncing ICCS parameters from database.") + logger.info(f"Syncing ICCS parameters from database for event {event.id}.") - active_event = get_active_event(session) - event_params = AimbatEventParametersBase.model_validate(active_event.parameters) + event_params = AimbatEventParametersBase.model_validate(event.parameters) for field_name in AimbatEventParametersBase.model_fields: if hasattr(iccs, field_name): setattr(iccs, field_name, getattr(event_params, field_name)) @@ -133,22 +163,24 @@ def run_iccs(session: Session, iccs: ICCS, autoflip: bool, autoselect: bool) -> _write_back_seismograms(session, iccs) -def run_mccc(session: Session, iccs: ICCS, all_seismograms: bool) -> None: +def run_mccc( + session: Session, event: AimbatEvent, iccs: ICCS, all_seismograms: bool +) -> None: """Run MCCC algorithm. Args: session: Database session. + event: AimbatEvent. iccs: ICCS instance. all_seismograms: Whether to include all seismograms in the MCCC processing, or just the selected ones. """ - logger.info(f"Running MCCC with {all_seismograms=}.") + logger.info(f"Running MCCC for event {event.id} with {all_seismograms=}.") - active_event = get_active_event(session) iccs.run_mccc( all_seismograms=all_seismograms, - min_cc=active_event.parameters.mccc_min_ccnorm, - damping=active_event.parameters.mccc_damp, + min_cc=event.parameters.mccc_min_ccnorm, + damping=event.parameters.mccc_damp, ) _write_back_seismograms(session, iccs) @@ -229,15 +261,18 @@ def update_pick( def update_timewindow( session: Session, + event: AimbatEvent, iccs: ICCS, context: bool, all: bool, use_seismogram_image: bool, return_fig: bool, ) -> tuple | None: - """Update the time window for the active event. + """Update the time window for the given event. Args: + session: Database session. + event: AimbatEvent. iccs: ICCS instance. context: Whether to use seismograms with extra context. all: Whether to plot all seismograms. @@ -248,16 +283,15 @@ def update_timewindow( A tuple of (Figure, Axes, widgets) if return_fig is True, otherwise None. """ - logger.info("Updating time window for active event.") + logger.info(f"Updating time window for event {event.id}.") result = _update_timewindow( # type: ignore[call-overload] iccs, context, all, use_seismogram_image, return_fig=return_fig ) if not return_fig: - active_event = get_active_event(session) - active_event.parameters.window_pre = iccs.window_pre - active_event.parameters.window_post = iccs.window_post + event.parameters.window_pre = iccs.window_pre + event.parameters.window_post = iccs.window_post session.commit() return None @@ -268,11 +302,18 @@ def update_timewindow( def update_min_ccnorm( - session: Session, iccs: ICCS, context: bool, all: bool, return_fig: bool + session: Session, + event: AimbatEvent, + iccs: ICCS, + context: bool, + all: bool, + return_fig: bool, ) -> tuple | None: - """Update the minimum cross correlation coefficient for the active event. + """Update the minimum cross correlation coefficient for the given event. Args: + session: Database session. + event: AimbatEvent. iccs: ICCS instance. context: Whether to use seismograms with extra context. all: Whether to plot all seismograms. @@ -282,13 +323,12 @@ def update_min_ccnorm( A tuple of (Figure, Axes, widgets) if return_fig is True, otherwise None. """ - logger.info("Updating minimum cross correlation coefficient for active event.") + logger.info(f"Updating minimum cross correlation coefficient for event {event.id}.") result = _update_min_ccnorm(iccs, context, all, return_fig=return_fig) # type: ignore[call-overload] if not return_fig: - active_event = get_active_event(session) - active_event.parameters.min_ccnorm = float(iccs.min_ccnorm) + event.parameters.min_ccnorm = float(iccs.min_ccnorm) session.commit() return None diff --git a/src/aimbat/core/_project.py b/src/aimbat/core/_project.py index 6e6436ee..e61590bb 100644 --- a/src/aimbat/core/_project.py +++ b/src/aimbat/core/_project.py @@ -1,22 +1,9 @@ -from aimbat.core import get_active_event -from aimbat.logger import logger -from aimbat.models import ( - AimbatEvent, - AimbatSeismogram, - AimbatStation, -) from sqlalchemy import Engine -from sqlalchemy.exc import NoResultFound -from sqlmodel import SQLModel, Session, select, text +from sqlmodel import SQLModel, text from pathlib import Path -from rich.console import Console -from rich.table import Table -from rich.panel import Panel -import aimbat.core._event as event -import aimbat.core._seismogram as seismogram -import aimbat.core._station as station +from aimbat.logger import logger -__all__ = ["create_project", "delete_project", "print_project_info"] +__all__ = ["create_project", "delete_project"] def _project_exists(engine: Engine) -> bool: @@ -86,11 +73,35 @@ def create_project(engine: Engine) -> None: BEFORE INSERT ON aimbatevent FOR EACH ROW WHEN NEW.active = TRUE BEGIN - UPDATE aimbatevent SET active = NULL + UPDATE aimbatevent SET active = NULL WHERE active = TRUE; END; """)) + # Trigger 3: Track last modification time when event parameters change + connection.execute(text(""" + CREATE TRIGGER IF NOT EXISTS event_modified_on_params_update + AFTER UPDATE ON aimbateventparameters + BEGIN + UPDATE aimbatevent SET last_modified = datetime('now') + WHERE id = NEW.event_id; + END; + """)) + + # Trigger 4: Track last modification time when seismogram parameters change + connection.execute(text(""" + CREATE TRIGGER IF NOT EXISTS event_modified_on_seis_params_update + AFTER UPDATE ON aimbatseismogramparameters + BEGIN + UPDATE aimbatevent + SET last_modified = strftime('%Y-%m-%d %H:%M:%f', 'now') + WHERE id = ( + SELECT event_id FROM aimbatseismogram + WHERE id = NEW.seismogram_id + ); + END; + """)) + def delete_project(engine: Engine) -> None: """Delete the AIMBAT project. @@ -114,74 +125,3 @@ def delete_project(engine: Engine) -> None: project_path.unlink() return raise RuntimeError("Unable to find/delete project.") - - -def print_project_info(engine: Engine) -> None: - """Show AIMBAT project information. - - Raises: - RuntimeError: If no project found. - """ - - logger.info("Printing project info.") - - if not _project_exists(engine): - raise RuntimeError( - 'No AIMBAT project found. Try running "aimbat project create" first.' - ) - - with Session(engine) as session: - grid = Table.grid(expand=False) - grid.add_column() - grid.add_column(justify="left") - if engine.driver == "pysqlite": - if engine.url.database == ":memory:": - grid.add_row("AIMBAT Project: ", "in-memory database") - else: - grid.add_row("AIMBAT Project File: ", str(engine.url.database)) - - events = len(session.exec(select(AimbatEvent)).all()) - completed_events = len(event.get_completed_events(session)) - stations = len(session.exec(select(AimbatStation)).all()) - seismograms = len(session.exec(select(AimbatSeismogram)).all()) - selected_seismograms = len( - seismogram.get_selected_seismograms(session, all_events=True) - ) - - grid.add_row( - "Number of Events (total/completed): ", - f"({events}/{completed_events})", - ) - - try: - active_event = get_active_event(session) - active_event_id = active_event.id - active_stations = len(station.get_stations_in_event(session, active_event)) - seismograms_in_event = len(active_event.seismograms) - selected_seismograms_in_event = len( - seismogram.get_selected_seismograms(session) - ) - except NoResultFound: - active_event_id = None - active_stations = None - seismograms_in_event = None - selected_seismograms_in_event = None - grid.add_row("Active Event ID: ", f"{active_event_id}") - grid.add_row( - "Number of Stations in Project (total/active event): ", - f"({stations}/{active_stations})", - ) - - grid.add_row( - "Number of Seismograms in Project (total/selected): ", - f"({seismograms}/{selected_seismograms})", - ) - grid.add_row( - "Number of Seismograms in Active Event (total/selected): ", - f"({seismograms_in_event}/{selected_seismograms_in_event})", - ) - - console = Console() - console.print( - Panel(grid, title="Project Info", title_align="left", border_style="dim") - ) diff --git a/src/aimbat/core/_seismogram.py b/src/aimbat/core/_seismogram.py index ba79d56e..ce66b20c 100644 --- a/src/aimbat/core/_seismogram.py +++ b/src/aimbat/core/_seismogram.py @@ -1,15 +1,18 @@ -import aimbat.core._event as event import uuid import matplotlib.pyplot as plt import matplotlib.dates as mdates -from aimbat.core import get_active_event +from pandas import Timestamp +from sqlmodel import Session, select +from sqlalchemy.exc import NoResultFound +from typing import overload +from collections.abc import Sequence +from pydantic import TypeAdapter +from typing import Any, Literal +from pysmo import MiniSeismogram +from pysmo.functions import detrend, normalize, clone_to_mini +from pysmo.tools.plotutils import time_array +from pysmo.tools.azdist import distance from aimbat.logger import logger -from aimbat.utils import ( - uuid_shortener, - make_table, - TABLE_STYLING, - json_to_table, -) from aimbat.models import ( AimbatEvent, AimbatSeismogram, @@ -21,18 +24,6 @@ SeismogramParameterBool, SeismogramParameterTimestamp, ) -from pysmo import MiniSeismogram -from pysmo.functions import detrend, normalize, clone_to_mini -from pysmo.tools.plotutils import time_array -from pysmo.tools.azdist import distance -from pandas import Timestamp -from rich.console import Console -from sqlmodel import Session, select -from sqlalchemy.exc import NoResultFound -from typing import overload -from collections.abc import Sequence -from pydantic import TypeAdapter -from typing import Any, Literal __all__ = [ "delete_seismogram_by_id", @@ -45,9 +36,7 @@ "reset_seismogram_parameters", "get_selected_seismograms", "dump_seismogram_table_to_json", - "print_seismogram_table", "dump_seismogram_parameter_table_to_json", - "print_seismogram_parameter_table", "plot_all_seismograms", ] @@ -274,12 +263,13 @@ def set_seismogram_parameter( def get_selected_seismograms( - session: Session, all_events: bool = False + session: Session, event: AimbatEvent | None = None, all_events: bool = False ) -> Sequence[AimbatSeismogram]: - """Get the selected seismograms for the active avent. + """Get the selected seismograms for the given event. Args: session: Database session. + event: Event to return selected seismograms for. all_events: Get the selected seismograms for all events. Returns: Selected seismograms. @@ -295,13 +285,14 @@ def get_selected_seismograms( .where(AimbatSeismogramParameters.select == 1) ) else: - logger.debug("Selecting seismograms for active event only.") + if event is None: + raise ValueError("An event must be provided when all_events is False.") + logger.debug(f"Selecting seismograms for event {event.id} only.") statement = ( select(AimbatSeismogram) .join(AimbatSeismogramParameters) - .join(AimbatEvent) .where(AimbatSeismogramParameters.select == 1) - .where(AimbatEvent.active == 1) + .where(AimbatSeismogram.event_id == event.id) ) seismograms = session.exec(statement).all() @@ -325,20 +316,36 @@ def dump_seismogram_table_to_json(session: Session) -> str: @overload def dump_seismogram_parameter_table_to_json( - session: Session, all_events: bool, as_string: Literal[True] + session: Session, + all_events: bool, + as_string: Literal[True], + event: AimbatEvent | None = None, ) -> str: ... @overload def dump_seismogram_parameter_table_to_json( - session: Session, all_events: bool, as_string: Literal[False] + session: Session, + all_events: bool, + as_string: Literal[False], + event: AimbatEvent | None = None, ) -> list[dict[str, Any]]: ... def dump_seismogram_parameter_table_to_json( - session: Session, all_events: bool, as_string: bool + session: Session, + all_events: bool, + as_string: bool, + event: AimbatEvent | None = None, ) -> str | list[dict[str, Any]]: - """Dump the seismogram parameter table data to json.""" + """Dump the seismogram parameter table data to json. + + Args: + session: Database session. + all_events: Include parameters for all events. + as_string: Return as JSON string. + event: Event to dump parameters for (only used when all_events is False). + """ logger.info("Dumping AimbatSeismogramParameters table to json.") @@ -349,11 +356,12 @@ def dump_seismogram_parameter_table_to_json( if all_events: parameters = session.exec(select(AimbatSeismogramParameters)).all() else: + if event is None: + raise ValueError("An event must be provided when all_events is False.") parameters = session.exec( select(AimbatSeismogramParameters) .join(AimbatSeismogram) - .join(AimbatEvent) - .where(AimbatEvent.active == 1) + .where(AimbatSeismogram.event_id == event.id) ).all() if as_string: @@ -361,150 +369,34 @@ def dump_seismogram_parameter_table_to_json( return adapter.dump_python(parameters, mode="json") -def print_seismogram_table( - session: Session, short: bool, all_events: bool = False -) -> None: - """Prints a pretty table with AIMBAT seismograms. - - Args: - short: Shorten and format the output to be more human-readable. - all_events: Print seismograms for all events. - """ - - logger.info("Printing AIMBAT seismogram table.") - - title = "AIMBAT seismograms for all events" - seismograms = None - - if all_events: - logger.debug("Selecting seismograms for all events.") - seismograms = session.exec(select(AimbatSeismogram)).all() - else: - logger.debug("Selecting seismograms for active event only.") - active_event = get_active_event(session) - seismograms = active_event.seismograms - if short: - title = f"AIMBAT seismograms for event {active_event.time.strftime('%Y-%m-%d %H:%M:%S')} (ID={event.uuid_shortener(session, active_event)})" - else: - title = f"AIMBAT seismograms for event {active_event.time} (ID={active_event.id})" - - logger.debug(f"Found {len(seismograms)} seismograms for the table.") - - table = make_table(title=title) - table.add_column( - "ID (shortened)" if short else "ID", - justify="center", - style=TABLE_STYLING.id, - no_wrap=True, - ) - table.add_column( - "Selected", justify="center", style=TABLE_STYLING.mine, no_wrap=True - ) - table.add_column("NPTS", justify="center", style=TABLE_STYLING.mine, no_wrap=True) - table.add_column("Delta", justify="center", style=TABLE_STYLING.mine, no_wrap=True) - table.add_column( - "Data ID", justify="center", style=TABLE_STYLING.linked, no_wrap=True - ) - table.add_column("Station ID", justify="center", style=TABLE_STYLING.linked) - table.add_column("Station Name", justify="center", style=TABLE_STYLING.linked) - if all_events: - table.add_column("Event ID", justify="center", style=TABLE_STYLING.linked) - - for seismogram in seismograms: - logger.debug(f"Adding seismogram with ID {seismogram.id} to the table.") - row = [ - (uuid_shortener(session, seismogram) if short else str(seismogram.id)), - TABLE_STYLING.bool_formatter(seismogram.parameters.select), - str(len(seismogram.data)), - str(seismogram.delta.total_seconds()), - ( - uuid_shortener(session, seismogram.datasource) - if short - else str(seismogram.datasource.id) - ), - ( - uuid_shortener(session, seismogram.station) - if short - else str(seismogram.station.id) - ), - f"{seismogram.station.name} - {seismogram.station.network}", - ] - - if all_events: - row.append( - uuid_shortener(session, seismogram.event) - if short - else str(seismogram.event.id) - ) - table.add_row(*row) - - console = Console() - console.print(table) - - -def print_seismogram_parameter_table(session: Session, short: bool) -> None: - """Print a pretty table with AIMBAT seismogram parameter values for the active event. - - Args: - short: Shorten and format the output to be more human-readable. - """ - - logger.info("Printing AIMBAT seismogram parameters table for active event.") - - active_event = get_active_event(session) - title = f"Seismogram parameters for event: {uuid_shortener(session, active_event) if short else str(active_event.id)}" - - json_to_table( - data=dump_seismogram_parameter_table_to_json( - session, all_events=False, as_string=False - ), - title=title, - skip_keys=["id"], - column_order=["seismogram_id", "select"], - common_column_kwargs={"highlight": True}, - formatters={ - "seismogram_id": lambda x: ( - uuid_shortener(session, AimbatSeismogram, str_uuid=x) if short else x - ), - }, - column_kwargs={ - "seismogram_id": { - "header": "Seismogram ID (shortened)" if short else "Seismogram ID", - "justify": "center", - "style": TABLE_STYLING.mine, - }, - }, - ) - - @overload def plot_all_seismograms( - session: Session, return_fig: Literal[True] + session: Session, event: AimbatEvent, return_fig: Literal[True] ) -> tuple[plt.Figure, plt.Axes]: ... @overload -def plot_all_seismograms(session: Session, return_fig: Literal[False]) -> None: ... +def plot_all_seismograms( + session: Session, event: AimbatEvent, return_fig: Literal[False] +) -> None: ... def plot_all_seismograms( - session: Session, return_fig: bool + session: Session, event: AimbatEvent, return_fig: bool ) -> tuple[plt.Figure, plt.Axes] | None: """Plot all seismograms for a particular event ordered by great circle distance. Args: session: Database session. + event: AimbatEvent. return_fig: Whether to return the figure and axes objects instead of showing the plot. Returns: figure and axes objects if return_fig is True, otherwise None. """ - if (active_event := get_active_event(session)) is None: - raise RuntimeError("No active event set.") - - if len(seismograms := active_event.seismograms) == 0: - raise RuntimeError("No seismograms found in active event.") + if len(seismograms := event.seismograms) == 0: + raise RuntimeError(f"No seismograms found in event {event.id}.") distance_dict = { seismogram.id: distance(seismogram.station, seismogram.event) / 1000 diff --git a/src/aimbat/core/_snapshot.py b/src/aimbat/core/_snapshot.py index c533d490..498bab88 100644 --- a/src/aimbat/core/_snapshot.py +++ b/src/aimbat/core/_snapshot.py @@ -1,8 +1,10 @@ import uuid import json -from aimbat.core import get_active_event +from sqlmodel import Session, select +from collections.abc import Sequence +from typing import overload, Literal, Any +from pydantic import TypeAdapter from aimbat.logger import logger -from aimbat.utils import uuid_shortener, json_to_table, TABLE_STYLING from aimbat.models import ( AimbatSnapshot, AimbatEvent, @@ -14,12 +16,6 @@ AimbatSeismogramParametersBase, AimbatEventParametersBase, ) -from sqlmodel import Session, select -from sqlalchemy import true -from pandas import Timestamp -from collections.abc import Sequence -from typing import overload, Literal, Any -from pydantic import TypeAdapter __all__ = [ "create_snapshot", @@ -29,30 +25,27 @@ "delete_snapshot", "get_snapshots", "dump_snapshot_tables_to_json", - "print_snapshot_table", - "print_snapshot_parameters_table_by_id", - "print_snapshot_parameters_table", ] -def create_snapshot(session: Session, comment: str | None = None) -> None: +def create_snapshot( + session: Session, event: AimbatEvent, comment: str | None = None +) -> None: """Create a snapshot of the AIMBAT processing parameters. Args: session: Database session. + event: AimbatEvent. comment: Optional comment. """ - active_aimbat_event = get_active_event(session) - logger.info( - f"Creating snapshot for event with id={active_aimbat_event.id} with {comment=}." - ) + logger.info(f"Creating snapshot for event with id={event.id} with {comment=}.") event_parameters_snapshot = AimbatEventParametersSnapshot.model_validate( - active_aimbat_event.parameters, + event.parameters, update={ "id": uuid.uuid4(), # we don't want to carry over the id from the active parameters - "parameters_id": active_aimbat_event.parameters.id, + "parameters_id": event.parameters.id, }, ) logger.debug( @@ -60,7 +53,7 @@ def create_snapshot(session: Session, comment: str | None = None) -> None: ) seismogram_parameter_snapshots = [] - for aimbat_seismogram in active_aimbat_event.seismograms: + for aimbat_seismogram in event.seismograms: seismogram_parameter_snapshot = AimbatSeismogramParametersSnapshot.model_validate( aimbat_seismogram.parameters, update={ @@ -74,7 +67,7 @@ def create_snapshot(session: Session, comment: str | None = None) -> None: seismogram_parameter_snapshots.append(seismogram_parameter_snapshot) aimbat_snapshot = AimbatSnapshot( - event=active_aimbat_event, + event=event, event_parameters_snapshot=event_parameters_snapshot, seismogram_parameters_snapshots=seismogram_parameter_snapshots, comment=comment, @@ -181,24 +174,26 @@ def delete_snapshot(session: Session, snapshot: AimbatSnapshot) -> None: def get_snapshots( - session: Session, all_events: bool = False + session: Session, event: AimbatEvent | None = None, all_events: bool = False ) -> Sequence[AimbatSnapshot]: - """Get the snapshots for the active avent. + """Get the snapshots for an event. Args: session: Database session. - all_events: Get the selected snapshots for all events. + event: Event to return snapshots for. + all_events: Get the snapshots for all events. Returns: Snapshots. """ logger.info("Getting AIMBAT snapshots.") - statement = ( - select(AimbatSnapshot) - .join(AimbatEvent) - .where(AimbatEvent.active == True if not all_events else true()) # noqa: E712 - ) + if all_events: + statement = select(AimbatSnapshot) + else: + if event is None: + raise ValueError("An event must be provided when all_events is False.") + statement = select(AimbatSnapshot).where(AimbatSnapshot.event_id == event.id) logger.debug(f"Executing statement to get snapshots: {statement}") return session.exec(statement).all() @@ -206,18 +201,27 @@ def get_snapshots( @overload def dump_snapshot_tables_to_json( - session: Session, all_events: bool, as_string: Literal[True] + session: Session, + all_events: bool, + as_string: Literal[True], + event: AimbatEvent | None = None, ) -> str: ... @overload def dump_snapshot_tables_to_json( - session: Session, all_events: bool, as_string: Literal[False] + session: Session, + all_events: bool, + as_string: Literal[False], + event: AimbatEvent | None = None, ) -> dict[str, list[dict[str, Any]]]: ... def dump_snapshot_tables_to_json( - session: Session, all_events: bool, as_string: bool + session: Session, + all_events: bool, + as_string: bool, + event: AimbatEvent | None = None, ) -> str | dict[str, list[dict[str, Any]]]: """Dump snapshot data as a dict of lists of dicts. @@ -233,10 +237,11 @@ def dump_snapshot_tables_to_json( session: Database session. all_events: Include snapshots for all events. as_string: Return a JSON string when True, otherwise a dict. + event: Event to dump snapshots for (only used when all_events is False). """ logger.info(f"Dumping AimbatSnapshot tables to json with {all_events=}.") - snapshots = get_snapshots(session, all_events) + snapshots = get_snapshots(session, event=event, all_events=all_events) snapshot_adapter: TypeAdapter[Sequence[_AimbatSnapshotRead]] = TypeAdapter( Sequence[_AimbatSnapshotRead] @@ -261,116 +266,3 @@ def dump_snapshot_tables_to_json( } return json.dumps(data) if as_string else data - - -def print_snapshot_table(session: Session, short: bool, all_events: bool) -> None: - """Print a pretty table with AIMBAT snapshots. - - Uses the ``snapshots`` portion of :func:`dump_snapshot_tables_to_json` - and renders it via :func:`~aimbat.utils.json_to_table`. - - Args: - session: Database session. - short: Shorten and format the output to be more human-readable. - all_events: Print all snapshots instead of limiting to the active event. - """ - - logger.info("Printing AIMBAT snapshots table.") - - title = "AIMBAT snapshots for all events" - - if not all_events: - active_event = get_active_event(session) - if short: - title = f"AIMBAT snapshots for event {active_event.time.strftime('%Y-%m-%d %H:%M:%S')} (ID={uuid_shortener(session, active_event)})" - else: - title = ( - f"AIMBAT snapshots for event {active_event.time} (ID={active_event.id})" - ) - - data = dump_snapshot_tables_to_json(session, all_events, as_string=False) - snapshot_data = data["snapshots"] - - column_order = ["id", "date", "comment", "seismogram_count"] - if all_events: - column_order.append("event_id") - - skip_keys = [] if all_events else ["event_id"] - - json_to_table( - data=snapshot_data, - title=title, - column_order=column_order, - skip_keys=skip_keys, - formatters={ - "id": lambda x: ( - uuid_shortener(session, AimbatSnapshot, str_uuid=x) if short else x - ), - "date": lambda x: TABLE_STYLING.timestamp_formatter(Timestamp(x), short), - "event_id": lambda x: ( - uuid_shortener(session, AimbatEvent, str_uuid=x) if short else x - ), - }, - common_column_kwargs={"justify": "center"}, - column_kwargs={ - "id": { - "header": "ID (shortened)" if short else "ID", - "style": TABLE_STYLING.id, - "no_wrap": True, - }, - "date": { - "header": "Date & Time", - "style": TABLE_STYLING.mine, - "no_wrap": True, - }, - "comment": {"style": TABLE_STYLING.mine}, - "seismogram_count": { - "header": "# Seismograms", - "style": TABLE_STYLING.linked, - }, - "selected_seismogram_count": { - "header": "# Selected", - "style": TABLE_STYLING.linked, - }, - "flipped_seismogram_count": { - "header": "# Flipped", - "style": TABLE_STYLING.linked, - }, - "event_id": { - "header": "Event ID (shortened)" if short else "Event ID", - "style": TABLE_STYLING.linked, - }, - }, - ) - - -def print_snapshot_parameters_table_by_id( - session: Session, snapshot_id: uuid.UUID, short: bool -) -> None: - """Print a pretty table with AIMBAT snapshot parameters for a given snapshot id.""" - snapshot = session.get(AimbatSnapshot, snapshot_id) - - if snapshot is None: - raise ValueError( - f"Unable to print snapshot parameters: snapshot with id={snapshot_id} not found." - ) - - print_snapshot_parameters_table(session, snapshot.event_parameters_snapshot, short) - - -def print_snapshot_parameters_table( - session: Session, snapshot: AimbatEventParametersSnapshot, short: bool -) -> None: - json_to_table( - data=snapshot.model_dump(mode="json"), - title=f"Saved event parameters in snapshot:: {uuid_shortener(session, snapshot.snapshot) if short else str(snapshot.snapshot.id)}", - skip_keys=["id", "snapshot_id", "parameters_id"], - common_column_kwargs={"highlight": True}, - column_kwargs={ - "Key": { - "header": "Parameter", - "justify": "left", - "style": TABLE_STYLING.id, - }, - }, - ) diff --git a/src/aimbat/core/_station.py b/src/aimbat/core/_station.py index 8739e712..ffe9bc47 100644 --- a/src/aimbat/core/_station.py +++ b/src/aimbat/core/_station.py @@ -1,23 +1,20 @@ import uuid -from aimbat.core import get_active_event -from aimbat.logger import logger -from aimbat.utils import uuid_shortener, json_to_table, TABLE_STYLING -from aimbat.models import AimbatStation, AimbatSeismogram, AimbatEvent from typing import overload, Literal, Any from sqlmodel import Session, select, col from sqlalchemy import func from sqlalchemy.exc import NoResultFound from collections.abc import Sequence from pydantic import TypeAdapter +from aimbat.logger import logger +from aimbat.models import AimbatStation, AimbatSeismogram, AimbatEvent __all__ = [ "delete_station_by_id", "delete_station", "get_stations_in_event", - "get_stations_in_active_event", - "get_stations_with_event_seismogram_count", "dump_station_table_to_json", - "print_station_table", + "dump_station_table_with_counts", + "get_stations_with_event_and_seismogram_count", ] @@ -55,35 +52,36 @@ def delete_station(session: Session, station: AimbatStation) -> None: @overload -def get_stations_in_active_event( - session: Session, as_json: Literal[False] +def get_stations_in_event( + session: Session, event: AimbatEvent, as_json: Literal[False] = ... ) -> Sequence[AimbatStation]: ... @overload -def get_stations_in_active_event( - session: Session, as_json: Literal[True] +def get_stations_in_event( + session: Session, event: AimbatEvent, as_json: Literal[True] ) -> list[dict[str, Any]]: ... -def get_stations_in_active_event( - session: Session, as_json: bool +def get_stations_in_event( + session: Session, event: AimbatEvent, as_json: bool = False ) -> Sequence[AimbatStation] | list[dict[str, Any]]: - """Get the stations for the active event. + """Get the stations for a particular event. Args: session: Database session. + event: Event to return stations for. + as_json: Whether to return the result as JSON. - Returns: Stations in active event. + Returns: Stations in event. """ - logger.info("Getting stations for active event.") + logger.info(f"Getting stations for event: {event.id}.") statement = ( select(AimbatStation) .distinct() .join(AimbatSeismogram) - .join(AimbatEvent) - .where(AimbatEvent.active == True) # noqa: E712 + .where(AimbatSeismogram.event_id == event.id) ) logger.debug(f"Executing query: {statement}") @@ -97,55 +95,16 @@ def get_stations_in_active_event( return adapter.dump_python(results, mode="json") -def get_stations_in_event( - session: Session, event: AimbatEvent -) -> Sequence[AimbatStation]: - """Get the stations for a particular event. - - Args: - session: Database session. - event: Event to return stations for. - - Returns: Stations in event. - """ - logger.info(f"Getting stations for event: {event.id}.") - - statement = ( - select(AimbatStation) - .join(AimbatSeismogram) - .join(AimbatEvent) - .where(AimbatEvent.id == event.id) - ) - - logger.debug(f"Executing query: {statement}") - stations = session.exec(statement).all() - - return stations - - -@overload -def get_stations_with_event_seismogram_count( - session: Session, as_json: Literal[False] -) -> Sequence[tuple[AimbatStation, int, int]]: ... - - -@overload -def get_stations_with_event_seismogram_count( - session: Session, as_json: Literal[True] -) -> list[dict[str, Any]]: ... - - -def get_stations_with_event_seismogram_count( - session: Session, as_json: bool -) -> Sequence[tuple[AimbatStation, int, int]] | list[dict[str, Any]]: +def get_stations_with_event_and_seismogram_count( + session: Session, +) -> Sequence[tuple[AimbatStation, int, int]]: """Get stations along with the count of seismograms and events they are associated with. Args: session: Database session. - as_json: Whether to return the result as JSON. Returns: A sequence of tuples containing the station, count of seismograms - and count of events, or a JSON string if as_json is True. + and count of events. """ logger.info("Getting stations with associated seismogram and event counts.") @@ -162,22 +121,27 @@ def get_stations_with_event_seismogram_count( ) logger.debug(f"Executing query: {statement}") - results = session.exec(statement).all() + return session.exec(statement).all() - if not as_json: - return results +def dump_station_table_with_counts(session: Session) -> list[dict[str, Any]]: + """Dump station table with associated seismogram and event counts to a list of dicts. + + Each dict represents a station and includes additional keys for the + seismogram and event counts. + + Args: + session: Database session. + + Returns: A list of dictionaries representing the stations with counts. + """ + results = get_stations_with_event_and_seismogram_count(session) formatted_results = [] for row in results: - # 1. Dump the station to a dict. mode="json" safely converts UUIDs/Datetimes to strings! station_dict = row[0].model_dump(mode="json") - - # 2. Add the counts directly to the dictionary station_dict["seismogram_count"] = row[1] station_dict["event_count"] = row[2] - - # 3. Add to our final list formatted_results.append(station_dict) return formatted_results @@ -191,119 +155,3 @@ def dump_station_table_to_json(session: Session) -> str: adapter: TypeAdapter[Sequence[AimbatStation]] = TypeAdapter(Sequence[AimbatStation]) aimbat_station = session.exec(select(AimbatStation)).all() return adapter.dump_json(aimbat_station).decode("utf-8") - - -def print_station_table( - session: Session, short: bool, all_events: bool = False -) -> None: - """Prints a pretty table with AIMBAT stations. - - Args: - session: Database session. - short: Shorten and format the output to be more human-readable. - all_events: Print stations for all events. - """ - logger.info("Printing station table.") - - title = "AIMBAT stations for all events" - - if all_events: - logger.debug("Selecting all AIMBAT stations.") - data = get_stations_with_event_seismogram_count(session, as_json=True) - else: - logger.debug("Selecting AIMBAT stations used by active event.") - active_event = get_active_event(session) - data = get_stations_in_active_event(session, as_json=True) - - if short: - title = f"AIMBAT stations for event {active_event.time.strftime('%Y-%m-%d %H:%M:%S')} (ID={uuid_shortener(session, active_event)})" - else: - title = ( - f"AIMBAT stations for event {active_event.time} (ID={active_event.id})" - ) - - column_order = [ - "id", - "name", - "network", - "channel", - "location", - "latitude", - "longitude", - "elevation", - ] - if all_events: - column_order.extend(["seismogram_count", "event_count"]) - - column_kwargs: dict[str, dict[str, Any]] = { - "id": { - "header": "ID (shortened)" if short else "ID", - "style": TABLE_STYLING.id, - "justify": "center", - "no_wrap": True, - }, - "name": { - "header": "Name", - "style": TABLE_STYLING.mine, - "justify": "center", - "no_wrap": True, - }, - "network": { - "header": "Network", - "style": TABLE_STYLING.mine, - "justify": "center", - "no_wrap": True, - }, - "channel": { - "header": "Channel", - "style": TABLE_STYLING.mine, - "justify": "center", - }, - "location": { - "header": "Location", - "style": TABLE_STYLING.mine, - "justify": "center", - }, - "latitude": { - "header": "Latitude", - "style": TABLE_STYLING.mine, - "justify": "center", - }, - "longitude": { - "header": "Longitude", - "style": TABLE_STYLING.mine, - "justify": "center", - }, - "elevation": { - "header": "Elevation", - "style": TABLE_STYLING.mine, - "justify": "center", - }, - "seismogram_count": { - "header": "# Seismograms", - "style": TABLE_STYLING.linked, - "justify": "center", - }, - "event_count": { - "header": "# Events", - "style": TABLE_STYLING.linked, - "justify": "center", - }, - } - - formatters = { - "id": lambda x: ( - uuid_shortener(session, AimbatStation, str_uuid=x) if short else str(x) - ), - "latitude": lambda x: f"{x:.3f}" if short else str(x), - "longitude": lambda x: f"{x:.3f}" if short else str(x), - "elevation": lambda x: f"{x:.0f}" if short else str(x), - } - - json_to_table( - data, - title=title, - column_order=column_order, - column_kwargs=column_kwargs, - formatters=formatters, - ) diff --git a/src/aimbat/models/__init__.py b/src/aimbat/models/__init__.py index 18653fa3..cd570c0b 100644 --- a/src/aimbat/models/__init__.py +++ b/src/aimbat/models/__init__.py @@ -30,6 +30,7 @@ from ._models import * from ._readers import * +from ._parameters import * __all__ = [s for s in dir() if not s.startswith("_") and s not in _internal_names] diff --git a/src/aimbat/models/_models.py b/src/aimbat/models/_models.py index 4fc7d811..80e76a0c 100644 --- a/src/aimbat/models/_models.py +++ b/src/aimbat/models/_models.py @@ -354,6 +354,11 @@ class AimbatEvent(SQLModel, table=True): latitude: float = Field(description="Event latitude.") longitude: float = Field(description="Event longitude.") depth: float | None = Field(default=None, description="Event depth.") + last_modified: PydanticTimestamp | None = Field( + default=None, + sa_type=SAPandasTimestamp, + description="Timestamp of the last modification to this event's parameters.", + ) seismograms: list[AimbatSeismogram] = Relationship( back_populates="event", cascade_delete=True ) diff --git a/src/aimbat/models/_parameters.py b/src/aimbat/models/_parameters.py index 6d1b8afc..269ce6d2 100644 --- a/src/aimbat/models/_parameters.py +++ b/src/aimbat/models/_parameters.py @@ -1,5 +1,8 @@ """Base classes defining AIMBAT processing parameters.""" +from sqlmodel import SQLModel, Field +from pydantic import model_validator +from typing import Self from aimbat import settings from aimbat._types import ( PydanticTimestamp, @@ -8,9 +11,6 @@ SAPandasTimestamp, SAPandasTimedelta, ) -from sqlmodel import SQLModel, Field -from pydantic import model_validator -from typing import Self __all__ = [ "AimbatEventParametersBase", diff --git a/src/aimbat/models/_readers.py b/src/aimbat/models/_readers.py index 64e11c29..c48f800c 100644 --- a/src/aimbat/models/_readers.py +++ b/src/aimbat/models/_readers.py @@ -18,6 +18,7 @@ class _AimbatEventRead(SQLModel): latitude: float longitude: float depth: float | None + last_modified: PydanticTimestamp | None = None completed: bool = False seismogram_count: int station_count: int @@ -32,6 +33,7 @@ def from_event(cls, event: "AimbatEvent") -> Self: latitude=event.latitude, longitude=event.longitude, depth=event.depth, + last_modified=event.last_modified, completed=event.parameters.completed, seismogram_count=event.seismogram_count, station_count=event.station_count, diff --git a/src/aimbat/utils/_style.py b/src/aimbat/utils/_style.py index bdaa3f66..f4f699e0 100644 --- a/src/aimbat/utils/_style.py +++ b/src/aimbat/utils/_style.py @@ -1,7 +1,7 @@ """AIMBAT styling.""" from dataclasses import dataclass -from pandas import Timestamp +from pandas import Timestamp, NaT from typing import Any from rich import box from rich.table import Table @@ -32,6 +32,8 @@ def bool_formatter(true_or_false: bool | Any) -> str: @staticmethod def timestamp_formatter(dt: Timestamp, short: bool) -> str: + if dt is NaT: + return "-" if short: return dt.strftime("%Y-%m-%d [light_sea_green]%H:%M:%S[/]") return str(dt) diff --git a/tests/integration/core/test_data.py b/tests/integration/core/test_data.py index 2dafbce6..bae713ae 100644 --- a/tests/integration/core/test_data.py +++ b/tests/integration/core/test_data.py @@ -12,9 +12,9 @@ from aimbat.io import DataType from aimbat.core import ( add_data_to_project, - get_data_for_active_event, - print_data_table, + get_data_for_event, dump_data_table_to_json, + get_active_event, ) from aimbat.models import ( AimbatDataSource, @@ -272,7 +272,8 @@ def test_get_data_sources_for_active_event(self, session: Session) -> None: Args: session (Session): Database session. """ - data_sources = get_data_for_active_event(session) + active_event = get_active_event(session) + data_sources = get_data_for_event(session, active_event) assert len(data_sources) != 0, "Expected data sources for the active event." assert all( isinstance(ds, AimbatDataSource) for ds in data_sources @@ -292,98 +293,6 @@ def test_dump_data_table_to_json(self, session: Session) -> None: returned_ids = [item["id"] for item in json_data] assert set(expected_ids) == set(returned_ids), "Expected IDs to match." - def test_print_data_table_for_all_events( - self, session: Session, capsys: pytest.CaptureFixture - ) -> None: - """Verifies that print_data_table produces output for all events. - - Args: - session (Session): Database session. - capsys (pytest.CaptureFixture): Fixture to capture stdout/stderr. - """ - print_data_table(session, short=False, all_events=True) - - expected_ids = session.exec(select(AimbatDataSource.id)).all() - - captured = capsys.readouterr() - assert "Data sources for all events" in captured.out - for id in expected_ids: - assert ( - str(id) in captured.out - ), "expected data source ID to be in the output table" - - def test_print_data_table_for_all_events_short( - self, session: Session, capsys: pytest.CaptureFixture - ) -> None: - """Verifies that print_data_table produces short output for all events. - - Args: - session (Session): Database session. - capsys (pytest.CaptureFixture): Fixture to capture stdout/stderr. - """ - expected_ids = session.exec(select(AimbatDataSource.id)).all() - - print_data_table(session, short=True, all_events=True) - - captured = capsys.readouterr() - assert "Data sources for all events" in captured.out - for id in expected_ids: - assert ( - str(id)[:2] in captured.out - ), "expected data source ID to be in the output table" - - def test_print_data_table_for_active_event( - self, session: Session, capsys: pytest.CaptureFixture - ) -> None: - """Verifies that print_data_table produces output for the active event. - - Args: - session (Session): Database session. - capsys (pytest.CaptureFixture): Fixture to capture stdout/stderr. - """ - statement = ( - select(AimbatDataSource.id) - .join(AimbatSeismogram) - .join(AimbatEvent) - .where(AimbatEvent.active == 1) - ) - expected_ids = session.exec(statement).all() - - print_data_table(session, short=False, all_events=False) - - captured = capsys.readouterr() - assert "Data sources for event" in captured.out - for id in expected_ids: - assert ( - str(id) in captured.out - ), "expected data source ID to be in the output table" - - def test_print_data_table_for_active_event_short( - self, session: Session, capsys: pytest.CaptureFixture - ) -> None: - """Verifies that print_data_table produces short output for the active event. - - Args: - session (Session): Database session. - capsys (pytest.CaptureFixture): Fixture to capture stdout/stderr. - """ - statement = ( - select(AimbatDataSource.id) - .join(AimbatSeismogram) - .join(AimbatEvent) - .where(AimbatEvent.active == 1) - ) - expected_ids = session.exec(statement).all() - - print_data_table(session, short=True, all_events=False) - - captured = capsys.readouterr() - assert "Data sources for event" in captured.out - for id in expected_ids: - assert ( - str(id)[:2] in captured.out - ), "expected data source ID to be in the output table" - # =================================================================== # Engine-level tests (add_data_to_project with engine fixture) diff --git a/tests/integration/core/test_event.py b/tests/integration/core/test_event.py index 1360eeab..736cc02c 100644 --- a/tests/integration/core/test_event.py +++ b/tests/integration/core/test_event.py @@ -4,8 +4,10 @@ import uuid import pytest from unittest.mock import patch -from aimbat.core import set_active_event, set_active_event_by_id, get_active_event -from aimbat.core._event import ( +from aimbat.core import ( + set_active_event, + set_active_event_by_id, + get_active_event, delete_event, delete_event_by_id, get_completed_events, @@ -14,8 +16,6 @@ set_event_parameter, dump_event_table_to_json, dump_event_parameter_table_to_json, - print_event_table, - print_event_parameter_table, ) from aimbat._types import EventParameter from aimbat.models import AimbatEvent, AimbatStation @@ -297,7 +297,8 @@ def test_get_timedelta_parameter(self, session: Session) -> None: Args: session: The database session. """ - value = get_event_parameter(session, EventParameter.WINDOW_PRE) + active_event = get_active_event(session) + value = get_event_parameter(session, active_event, EventParameter.WINDOW_PRE) assert isinstance(value, Timedelta) def test_get_float_parameter(self, session: Session) -> None: @@ -306,7 +307,8 @@ def test_get_float_parameter(self, session: Session) -> None: Args: session: The database session. """ - value = get_event_parameter(session, EventParameter.MIN_CCNORM) + active_event = get_active_event(session) + value = get_event_parameter(session, active_event, EventParameter.MIN_CCNORM) assert isinstance(value, float) def test_get_bool_parameter(self, session: Session) -> None: @@ -315,7 +317,8 @@ def test_get_bool_parameter(self, session: Session) -> None: Args: session: The database session. """ - value = get_event_parameter(session, EventParameter.COMPLETED) + active_event = get_active_event(session) + value = get_event_parameter(session, active_event, EventParameter.COMPLETED) assert isinstance(value, bool) @@ -328,9 +331,15 @@ def test_set_timedelta_parameter(self, session: Session) -> None: Args: session: The database session. """ + active_event = get_active_event(session) new_value = Timedelta(seconds=20) - set_event_parameter(session, EventParameter.WINDOW_POST, new_value) - assert get_event_parameter(session, EventParameter.WINDOW_POST) == new_value + set_event_parameter( + session, active_event, EventParameter.WINDOW_POST, new_value + ) + assert ( + get_event_parameter(session, active_event, EventParameter.WINDOW_POST) + == new_value + ) def test_set_float_parameter(self, session: Session) -> None: """Verifies that a float parameter is persisted correctly. @@ -338,9 +347,13 @@ def test_set_float_parameter(self, session: Session) -> None: Args: session: The database session. """ + active_event = get_active_event(session) new_value = 0.75 - set_event_parameter(session, EventParameter.MIN_CCNORM, new_value) - assert get_event_parameter(session, EventParameter.MIN_CCNORM) == new_value + set_event_parameter(session, active_event, EventParameter.MIN_CCNORM, new_value) + assert ( + get_event_parameter(session, active_event, EventParameter.MIN_CCNORM) + == new_value + ) def test_set_bool_parameter(self, session: Session) -> None: """Verifies that a bool parameter is persisted correctly. @@ -348,8 +361,11 @@ def test_set_bool_parameter(self, session: Session) -> None: Args: session: The database session. """ - set_event_parameter(session, EventParameter.COMPLETED, True) - assert get_event_parameter(session, EventParameter.COMPLETED) is True + active_event = get_active_event(session) + set_event_parameter(session, active_event, EventParameter.COMPLETED, True) + assert ( + get_event_parameter(session, active_event, EventParameter.COMPLETED) is True + ) # =================================================================== @@ -394,8 +410,9 @@ def test_active_event_as_string(self, session: Session) -> None: Args: session: The database session. """ + active_event = get_active_event(session) result = dump_event_parameter_table_to_json( - session, all_events=False, as_string=True + session, all_events=False, as_string=True, event=active_event ) assert isinstance(result, str) parsed = json.loads(result) @@ -409,8 +426,9 @@ def test_active_event_as_dict(self, session: Session) -> None: Args: session: The database session. """ + active_event = get_active_event(session) result = dump_event_parameter_table_to_json( - session, all_events=False, as_string=False + session, all_events=False, as_string=False, event=active_event ) assert isinstance(result, dict) assert "min_ccnorm" in result @@ -443,60 +461,3 @@ def test_all_events_as_list(self, session: Session) -> None: assert isinstance(result, list) assert len(result) > 0 assert "min_ccnorm" in result[0] - - -# =================================================================== -# Print tables -# =================================================================== - - -class TestPrintEventTable: - """Tests for printing the event table.""" - - def test_print_short(self, session: Session, capsys: pytest.CaptureFixture) -> None: - """Verifies that print_event_table produces output with short=True. - - Args: - session: The database session. - capsys: The pytest capsys fixture. - """ - print_event_table(session, short=True) - assert len(capsys.readouterr().out) > 0 - - def test_print_long(self, session: Session, capsys: pytest.CaptureFixture) -> None: - """Verifies that print_event_table produces output with short=False. - - Args: - session: The database session. - capsys: The pytest capsys fixture. - """ - print_event_table(session, short=False) - assert len(capsys.readouterr().out) > 0 - - -class TestPrintEventParameterTable: - """Tests for printing the event parameter table.""" - - def test_print_active_event( - self, session: Session, capsys: pytest.CaptureFixture - ) -> None: - """Verifies that print_event_parameter_table produces output for the active event. - - Args: - session: The database session. - capsys: The pytest capsys fixture. - """ - print_event_parameter_table(session, short=False, all_events=False) - assert len(capsys.readouterr().out) > 0 - - def test_print_all_events( - self, session: Session, capsys: pytest.CaptureFixture - ) -> None: - """Verifies that print_event_parameter_table produces output for all events. - - Args: - session: The database session. - capsys: The pytest capsys fixture. - """ - print_event_parameter_table(session, short=False, all_events=True) - assert len(capsys.readouterr().out) > 0 diff --git a/tests/integration/core/test_project.py b/tests/integration/core/test_project.py index 43f63f68..54cefc22 100644 --- a/tests/integration/core/test_project.py +++ b/tests/integration/core/test_project.py @@ -6,7 +6,7 @@ import pytest from pathlib import Path from aimbat.core import create_project, delete_project -from aimbat.core._project import _project_exists, print_project_info +from aimbat.core._project import _project_exists from collections.abc import Generator from sqlalchemy import Engine @@ -104,28 +104,4 @@ def test_raises_when_no_project( capsys: The pytest capsys fixture. """ with pytest.raises(RuntimeError): - print_project_info(engine_from_file) - - def test_with_empty_project( - self, patched_engine: Engine, capsys: pytest.CaptureFixture - ) -> None: - """Verifies that output is produced for a project with no data or active event. - - Args: - patched_engine: The monkeypatched SQLAlchemy Engine. - capsys: The pytest capsys fixture. - """ - print_project_info(patched_engine) - assert len(capsys.readouterr().out) > 0 - - def test_with_data_and_active_event( - self, loaded_engine: Engine, capsys: pytest.CaptureFixture - ) -> None: - """Verifies that output is produced for a project with data and an active event. - - Args: - loaded_engine: The monkeypatched SQLAlchemy Engine with data loaded. - capsys: The pytest capsys fixture. - """ - print_project_info(loaded_engine) - assert len(capsys.readouterr().out) > 0 + delete_project(engine_from_file) diff --git a/tests/integration/core/test_seismogram.py b/tests/integration/core/test_seismogram.py index e2856308..c3bf2fda 100644 --- a/tests/integration/core/test_seismogram.py +++ b/tests/integration/core/test_seismogram.py @@ -3,7 +3,7 @@ import json import uuid import pytest -from aimbat.core._seismogram import ( +from aimbat.core import ( delete_seismogram, delete_seismogram_by_id, get_seismogram_parameter, @@ -15,9 +15,8 @@ get_selected_seismograms, dump_seismogram_table_to_json, dump_seismogram_parameter_table_to_json, - print_seismogram_table, - print_seismogram_parameter_table, plot_all_seismograms, + get_active_event, ) from aimbat.models._parameters import AimbatSeismogramParametersBase from aimbat._types import SeismogramParameter @@ -276,7 +275,8 @@ def test_all_selected_by_default(self, session: Session) -> None: Args: session: The database session. """ - selected = get_selected_seismograms(session) + active_event = get_active_event(session) + selected = get_selected_seismograms(session, event=active_event) assert len(selected) > 0 def test_after_deselecting_one( @@ -288,9 +288,13 @@ def test_after_deselecting_one( session: The database session. seismogram: An AimbatSeismogram to deselect. """ - count_before = len(get_selected_seismograms(session)) + active_event = get_active_event(session) + count_before = len(get_selected_seismograms(session, event=active_event)) set_seismogram_parameter(session, seismogram, SeismogramParameter.SELECT, False) - assert len(get_selected_seismograms(session)) == count_before - 1 + assert ( + len(get_selected_seismograms(session, event=active_event)) + == count_before - 1 + ) def test_all_events(self, session: Session) -> None: """Verifies that get_selected_seismograms returns seismograms across all events. @@ -298,7 +302,10 @@ def test_all_events(self, session: Session) -> None: Args: session: The database session. """ - selected_active = get_selected_seismograms(session, all_events=False) + active_event = get_active_event(session) + selected_active = get_selected_seismograms( + session, event=active_event, all_events=False + ) selected_all = get_selected_seismograms(session, all_events=True) assert len(selected_all) >= len(selected_active) @@ -328,8 +335,9 @@ def test_active_event_as_string(self, session: Session) -> None: Args: session: The database session. """ + active_event = get_active_event(session) result = dump_seismogram_parameter_table_to_json( - session, all_events=False, as_string=True + session, all_events=False, as_string=True, event=active_event ) assert isinstance(result, str) parsed = json.loads(result) @@ -342,8 +350,9 @@ def test_active_event_as_list(self, session: Session) -> None: Args: session: The database session. """ + active_event = get_active_event(session) result = dump_seismogram_parameter_table_to_json( - session, all_events=False, as_string=False + session, all_events=False, as_string=False, event=active_event ) assert isinstance(result, list) assert len(result) > 0 @@ -382,8 +391,9 @@ def test_all_events_returns_more_than_active_only(self, session: Session) -> Non Args: session: The database session. """ + active_event = get_active_event(session) active_only = dump_seismogram_parameter_table_to_json( - session, all_events=False, as_string=False + session, all_events=False, as_string=False, event=active_event ) all_events = dump_seismogram_parameter_table_to_json( session, all_events=True, as_string=False @@ -391,68 +401,6 @@ def test_all_events_returns_more_than_active_only(self, session: Session) -> Non assert len(all_events) >= len(active_only) -class TestPrintSeismogramTable: - """Tests for printing the seismogram table.""" - - def test_active_event_short( - self, session: Session, capsys: pytest.CaptureFixture - ) -> None: - """Verifies that output is produced for the active event with short=True. - - Args: - session: The database session. - capsys: The pytest capsys fixture. - """ - print_seismogram_table(session, short=True, all_events=False) - assert len(capsys.readouterr().out) > 0 - - def test_active_event_long( - self, session: Session, capsys: pytest.CaptureFixture - ) -> None: - """Verifies that output is produced for the active event with short=False. - - Args: - session: The database session. - capsys: The pytest capsys fixture. - """ - print_seismogram_table(session, short=False, all_events=False) - assert len(capsys.readouterr().out) > 0 - - def test_all_events(self, session: Session, capsys: pytest.CaptureFixture) -> None: - """Verifies that output is produced when printing seismograms for all events. - - Args: - session: The database session. - capsys: The pytest capsys fixture. - """ - print_seismogram_table(session, short=False, all_events=True) - assert len(capsys.readouterr().out) > 0 - - -class TestPrintSeismogramParameterTable: - """Tests for printing the seismogram parameter table.""" - - def test_print_short(self, session: Session, capsys: pytest.CaptureFixture) -> None: - """Verifies that output is produced with short=True. - - Args: - session: The database session. - capsys: The pytest capsys fixture. - """ - print_seismogram_parameter_table(session, short=True) - assert len(capsys.readouterr().out) > 0 - - def test_print_long(self, session: Session, capsys: pytest.CaptureFixture) -> None: - """Verifies that output is produced with short=False. - - Args: - session: The database session. - capsys: The pytest capsys fixture. - """ - print_seismogram_parameter_table(session, short=False) - assert len(capsys.readouterr().out) > 0 - - class TestPlotAllSeismograms: """Tests for plotting seismograms.""" @@ -462,5 +410,6 @@ def test_returns_figure(self, session: Session) -> None: Args: session: The database session. """ - fig, _ = plot_all_seismograms(session, return_fig=True) + active_event = get_active_event(session) + fig, _ = plot_all_seismograms(session, event=active_event, return_fig=True) assert isinstance(fig, Figure) diff --git a/tests/integration/core/test_snapshots.py b/tests/integration/core/test_snapshots.py index 64623e20..895e952b 100644 --- a/tests/integration/core/test_snapshots.py +++ b/tests/integration/core/test_snapshots.py @@ -12,9 +12,6 @@ rollback_to_snapshot, rollback_to_snapshot_by_id, dump_snapshot_tables_to_json, - print_snapshot_table, - print_snapshot_parameters_table, - print_snapshot_parameters_table_by_id, ) from aimbat.core import get_active_event from aimbat.models import AimbatSnapshot, AimbatSeismogram @@ -44,7 +41,8 @@ def snapshot(session: Session) -> AimbatSnapshot: Returns: An AimbatSnapshot for the active event. """ - create_snapshot(session) + active_event = get_active_event(session) + create_snapshot(session, active_event) return session.exec(select(AimbatSnapshot)).one() @@ -58,7 +56,8 @@ def test_creates_snapshot(self, session: Session) -> None: session: The database session. """ assert len(session.exec(select(AimbatSnapshot)).all()) == 0 - create_snapshot(session) + active_event = get_active_event(session) + create_snapshot(session, active_event) assert len(session.exec(select(AimbatSnapshot)).all()) == 1 def test_snapshot_linked_to_active_event(self, session: Session) -> None: @@ -68,7 +67,7 @@ def test_snapshot_linked_to_active_event(self, session: Session) -> None: session: The database session. """ active_event = get_active_event(session) - create_snapshot(session) + create_snapshot(session, active_event) snapshot = session.exec(select(AimbatSnapshot)).one() assert snapshot.event_id == active_event.id @@ -78,7 +77,8 @@ def test_snapshot_with_comment(self, session: Session) -> None: Args: session: The database session. """ - create_snapshot(session, comment="test comment") + active_event = get_active_event(session) + create_snapshot(session, active_event, comment="test comment") snapshot = session.exec(select(AimbatSnapshot)).one() assert snapshot.comment == "test comment" @@ -88,7 +88,8 @@ def test_snapshot_without_comment(self, session: Session) -> None: Args: session: The database session. """ - create_snapshot(session) + active_event = get_active_event(session) + create_snapshot(session, active_event) snapshot = session.exec(select(AimbatSnapshot)).one() assert snapshot.comment is None @@ -101,7 +102,7 @@ def test_snapshot_captures_seismogram_parameters(self, session: Session) -> None active_event = get_active_event(session) n_seismograms = len(active_event.seismograms) - create_snapshot(session) + create_snapshot(session, active_event) snapshot = session.exec(select(AimbatSnapshot)).one() assert len(snapshot.seismogram_parameters_snapshots) == n_seismograms @@ -311,7 +312,8 @@ def test_no_snapshots_initially(self, session: Session) -> None: Args: session: The database session. """ - assert len(get_snapshots(session)) == 0 + active_event = get_active_event(session) + assert len(get_snapshots(session, event=active_event)) == 0 def test_get_snapshots_for_active_event( self, session: Session, snapshot: AimbatSnapshot @@ -322,7 +324,8 @@ def test_get_snapshots_for_active_event( session: The database session. snapshot: An AimbatSnapshot for the active event. """ - snapshots = get_snapshots(session, all_events=False) + active_event = get_active_event(session) + snapshots = get_snapshots(session, event=active_event, all_events=False) assert len(snapshots) == 1 assert snapshots[0].id == snapshot.id @@ -344,9 +347,10 @@ def test_multiple_snapshots(self, session: Session) -> None: Args: session: The database session. """ - create_snapshot(session, comment="first") - create_snapshot(session, comment="second") - assert len(get_snapshots(session)) == 2 + active_event = get_active_event(session) + create_snapshot(session, active_event, comment="first") + create_snapshot(session, active_event, comment="second") + assert len(get_snapshots(session, event=active_event)) == 2 class TestDumpSnapshotTablesToJson: @@ -359,7 +363,10 @@ def test_as_string(self, session: Session, snapshot: AimbatSnapshot) -> None: session: The database session. snapshot: An AimbatSnapshot to include in the dump. """ - result = dump_snapshot_tables_to_json(session, all_events=False, as_string=True) + active_event = get_active_event(session) + result = dump_snapshot_tables_to_json( + session, all_events=False, as_string=True, event=active_event + ) assert isinstance(result, str) parsed = json.loads(result) assert "snapshots" in parsed @@ -373,8 +380,9 @@ def test_as_dict(self, session: Session, snapshot: AimbatSnapshot) -> None: session: The database session. snapshot: An AimbatSnapshot to include in the dump. """ + active_event = get_active_event(session) result = dump_snapshot_tables_to_json( - session, all_events=False, as_string=False + session, all_events=False, as_string=False, event=active_event ) assert isinstance(result, dict) assert "snapshots" in result @@ -389,8 +397,9 @@ def test_all_events_includes_more_snapshots( session: The database session. snapshot: An AimbatSnapshot to include in the dump. """ + active_event = get_active_event(session) active_only = dump_snapshot_tables_to_json( - session, all_events=False, as_string=False + session, all_events=False, as_string=False, event=active_event ) all_events = dump_snapshot_tables_to_json( session, all_events=True, as_string=False @@ -409,136 +418,3 @@ def test_seismogram_parameters_count( n_seismograms = len(session.exec(select(AimbatSeismogram)).all()) result = dump_snapshot_tables_to_json(session, all_events=True, as_string=False) assert len(result["seismogram_parameters"]) <= n_seismograms - - -class TestPrintSnapshotParametersTable: - """Tests for printing the snapshot parameters table.""" - - def test_print_short( - self, - session: Session, - snapshot: AimbatSnapshot, - capsys: pytest.CaptureFixture, - ) -> None: - """Verifies that output is produced with short=True. - - Args: - session: The database session. - snapshot: An AimbatSnapshot whose event parameters are displayed. - capsys: The pytest capsys fixture. - """ - print_snapshot_parameters_table( - session, snapshot.event_parameters_snapshot, short=True - ) - assert len(capsys.readouterr().out) > 0 - - def test_print_long( - self, - session: Session, - snapshot: AimbatSnapshot, - capsys: pytest.CaptureFixture, - ) -> None: - """Verifies that output is produced with short=False. - - Args: - session: The database session. - snapshot: An AimbatSnapshot whose event parameters are displayed. - capsys: The pytest capsys fixture. - """ - print_snapshot_parameters_table( - session, snapshot.event_parameters_snapshot, short=False - ) - assert len(capsys.readouterr().out) > 0 - - def test_print_by_id_short( - self, - session: Session, - snapshot: AimbatSnapshot, - capsys: pytest.CaptureFixture, - ) -> None: - """Verifies that output is produced when looking up by snapshot ID with short=True. - - Args: - session: The database session. - snapshot: An AimbatSnapshot whose ID is used for lookup. - capsys: The pytest capsys fixture. - """ - print_snapshot_parameters_table_by_id(session, snapshot.id, short=True) - assert len(capsys.readouterr().out) > 0 - - def test_print_by_id_long( - self, - session: Session, - snapshot: AimbatSnapshot, - capsys: pytest.CaptureFixture, - ) -> None: - """Verifies that output is produced when looking up by snapshot ID with short=False. - - Args: - session: The database session. - snapshot: An AimbatSnapshot whose ID is used for lookup. - capsys: The pytest capsys fixture. - """ - print_snapshot_parameters_table_by_id(session, snapshot.id, short=False) - assert len(capsys.readouterr().out) > 0 - - def test_print_by_id_not_found(self, session: Session) -> None: - """Verifies that a ValueError is raised for an unknown snapshot ID. - - Args: - session: The database session. - """ - with pytest.raises(ValueError): - print_snapshot_parameters_table_by_id(session, uuid.uuid4(), short=False) - - -class TestPrintSnapshotTable: - """Tests for printing the snapshot table.""" - - def test_print_active_event_short( - self, - session: Session, - snapshot: AimbatSnapshot, - capsys: pytest.CaptureFixture, - ) -> None: - """Verifies that output is produced for the active event with short=True. - - Args: - session: The database session. - snapshot: An AimbatSnapshot to display. - capsys: The pytest capsys fixture. - """ - print_snapshot_table(session, short=True, all_events=False) - assert len(capsys.readouterr().out) > 0 - - def test_print_active_event_long( - self, - session: Session, - snapshot: AimbatSnapshot, - capsys: pytest.CaptureFixture, - ) -> None: - """Verifies that output is produced for the active event with short=False. - - Args: - session: The database session. - snapshot: An AimbatSnapshot to display. - capsys: The pytest capsys fixture. - """ - print_snapshot_table(session, short=False, all_events=False) - assert len(capsys.readouterr().out) > 0 - - def test_print_all_events( - self, - session: Session, - snapshot: AimbatSnapshot, - capsys: pytest.CaptureFixture, - ) -> None: - """Verifies that output is produced when printing snapshots for all events. - - Args: - session: The database session. - snapshot: An AimbatSnapshot to display. - capsys: The pytest capsys fixture. - """ - print_snapshot_table(session, short=False, all_events=True) - assert len(capsys.readouterr().out) > 0 diff --git a/tests/integration/core/test_station.py b/tests/integration/core/test_station.py index 74f69b9f..6e10f936 100644 --- a/tests/integration/core/test_station.py +++ b/tests/integration/core/test_station.py @@ -11,10 +11,9 @@ delete_station, delete_station_by_id, dump_station_table_to_json, - get_stations_in_active_event, get_stations_in_event, - get_stations_with_event_seismogram_count, - print_station_table, + get_stations_with_event_and_seismogram_count, + dump_station_table_with_counts, ) from aimbat.models import AimbatStation @@ -96,7 +95,8 @@ def test_returns_stations(self, session: Session) -> None: Args: session: The database session. """ - stations = get_stations_in_active_event(session, as_json=False) + active_event = get_active_event(session) + stations = get_stations_in_event(session, active_event, as_json=False) assert len(stations) > 0, "Expected at least one station for the active event" def test_returns_aimbat_station_instances(self, session: Session) -> None: @@ -105,7 +105,8 @@ def test_returns_aimbat_station_instances(self, session: Session) -> None: Args: session: The database session. """ - stations = get_stations_in_active_event(session, as_json=False) + active_event = get_active_event(session) + stations = get_stations_in_event(session, active_event, as_json=False) assert all( isinstance(s, AimbatStation) for s in stations ), "All returned items should be AimbatStation instances" @@ -116,7 +117,8 @@ def test_as_json_returns_list_of_dicts(self, session: Session) -> None: Args: session: The database session. """ - result = get_stations_in_active_event(session, as_json=True) + active_event = get_active_event(session) + result = get_stations_in_event(session, active_event, as_json=True) assert isinstance(result, list), "Expected a list when as_json=True" assert all( isinstance(item, dict) for item in result @@ -128,8 +130,9 @@ def test_as_json_count_matches_objects(self, session: Session) -> None: Args: session: The database session. """ - objects = get_stations_in_active_event(session, as_json=False) - json_list = get_stations_in_active_event(session, as_json=True) + active_event = get_active_event(session) + objects = get_stations_in_event(session, active_event, as_json=False) + json_list = get_stations_in_event(session, active_event, as_json=True) assert len(objects) == len( json_list ), "Object and JSON representations should have the same length" @@ -142,7 +145,7 @@ def test_stations_belong_to_active_event(self, session: Session) -> None: """ active_event = get_active_event(session) active_station_ids = {s.station_id for s in active_event.seismograms} - stations = get_stations_in_active_event(session, as_json=False) + stations = get_stations_in_event(session, active_event, as_json=False) returned_ids = {s.id for s in stations} assert ( returned_ids == active_station_ids @@ -198,7 +201,7 @@ def test_returns_all_stations(self, session: Session) -> None: session: The database session. """ all_stations = session.exec(select(AimbatStation)).all() - results = get_stations_with_event_seismogram_count(session, as_json=False) + results = get_stations_with_event_and_seismogram_count(session) assert len(results) == len( all_stations ), "Expected one row per station in the database" @@ -209,7 +212,7 @@ def test_returns_tuples_with_counts(self, session: Session) -> None: Args: session: The database session. """ - results = get_stations_with_event_seismogram_count(session, as_json=False) + results = get_stations_with_event_and_seismogram_count(session) for row in results: station, seismogram_count, event_count = row assert isinstance( @@ -228,7 +231,7 @@ def test_counts_are_non_negative(self, session: Session) -> None: Args: session: The database session. """ - results = get_stations_with_event_seismogram_count(session, as_json=False) + results = get_stations_with_event_and_seismogram_count(session) for _, seismogram_count, event_count in results: assert seismogram_count >= 0, "Seismogram count should be non-negative" assert event_count >= 0, "Event count should be non-negative" @@ -239,7 +242,7 @@ def test_as_json_returns_list_of_dicts(self, session: Session) -> None: Args: session: The database session. """ - results = get_stations_with_event_seismogram_count(session, as_json=True) + results = dump_station_table_with_counts(session) assert isinstance(results, list), "Expected a list when as_json=True" for item in results: assert isinstance(item, dict), "Each element should be a dict" @@ -252,8 +255,8 @@ def test_as_json_count_matches_objects(self, session: Session) -> None: Args: session: The database session. """ - objects = get_stations_with_event_seismogram_count(session, as_json=False) - json_list = get_stations_with_event_seismogram_count(session, as_json=True) + objects = get_stations_with_event_and_seismogram_count(session) + json_list = dump_station_table_with_counts(session) assert len(objects) == len( json_list ), "Object and JSON representations should have the same number of rows" @@ -294,49 +297,3 @@ def test_entries_contain_id_field(self, session: Session) -> None: result = json.loads(dump_station_table_to_json(session)) for entry in result: assert "id" in entry, "Each station entry should have an 'id' field" - - -class TestPrintStationTable: - """Tests for printing the station table.""" - - def test_print_active_event_short( - self, session: Session, capsys: pytest.CaptureFixture - ) -> None: - """Verifies that output is produced for the active event with short=True. - - Args: - session: The database session. - capsys: The pytest capsys fixture. - """ - print_station_table(session, short=True, all_events=False) - assert ( - len(capsys.readouterr().out) > 0 - ), "Expected output when printing station table (short, active event)" - - def test_print_active_event_long( - self, session: Session, capsys: pytest.CaptureFixture - ) -> None: - """Verifies that output is produced for the active event with short=False. - - Args: - session: The database session. - capsys: The pytest capsys fixture. - """ - print_station_table(session, short=False, all_events=False) - assert ( - len(capsys.readouterr().out) > 0 - ), "Expected output when printing station table (long, active event)" - - def test_print_all_events( - self, session: Session, capsys: pytest.CaptureFixture - ) -> None: - """Verifies that output is produced when printing stations for all events. - - Args: - session: The database session. - capsys: The pytest capsys fixture. - """ - print_station_table(session, short=False, all_events=True) - assert ( - len(capsys.readouterr().out) > 0 - ), "Expected output when printing station table for all events" diff --git a/tests/integration/models/test_models.py b/tests/integration/models/test_models.py index e93e1950..b6f90ebd 100644 --- a/tests/integration/models/test_models.py +++ b/tests/integration/models/test_models.py @@ -224,9 +224,10 @@ def test_delete_event_cascades_to_snapshots(self, session: Session) -> None: session.commit() # Create a snapshot via the core helper (uses the active event). - from aimbat.core import create_snapshot + from aimbat.core import create_snapshot, get_active_event - create_snapshot(session, comment="before delete") + active_event = get_active_event(session) + create_snapshot(session, active_event, comment="before delete") assert len(session.exec(select(AimbatSnapshot)).all()) == 1 assert len(session.exec(select(AimbatEventParametersSnapshot)).all()) == 1 assert len(session.exec(select(AimbatSeismogramParametersSnapshot)).all()) == 1 @@ -295,9 +296,10 @@ def test_delete_snapshot_cascades_to_parameter_snapshots( _make_seismogram(session, ev, sta) session.commit() - from aimbat.core import create_snapshot + from aimbat.core import create_snapshot, get_active_event - create_snapshot(session) + active_event = get_active_event(session) + create_snapshot(session, active_event) snapshot = session.exec(select(AimbatSnapshot)).one() session.delete(snapshot) diff --git a/tests/integration/models/test_operations.py b/tests/integration/models/test_operations.py index 33bfe15e..11c2fcf8 100644 --- a/tests/integration/models/test_operations.py +++ b/tests/integration/models/test_operations.py @@ -214,7 +214,8 @@ def test_snapshot_has_event_parameters_snapshot(self, session: Session) -> None: Args: session: The database session. """ - create_snapshot(session) + active_event = get_active_event(session) + create_snapshot(session, active_event) snapshot = session.exec(select(AimbatSnapshot)).one() assert isinstance( snapshot.event_parameters_snapshot, AimbatEventParametersSnapshot @@ -228,7 +229,8 @@ def test_snapshot_has_seismogram_parameter_snapshots( Args: session: The database session. """ - create_snapshot(session) + active_event = get_active_event(session) + create_snapshot(session, active_event) snapshot = session.exec(select(AimbatSnapshot)).one() assert len(snapshot.seismogram_parameters_snapshots) > 0 assert all( @@ -242,7 +244,8 @@ def test_snapshot_back_reference_to_event(self, session: Session) -> None: Args: session: The database session. """ - create_snapshot(session) + active_event = get_active_event(session) + create_snapshot(session, active_event) snapshot = session.exec(select(AimbatSnapshot)).one() assert isinstance(snapshot.event, AimbatEvent) @@ -252,7 +255,8 @@ def test_snapshot_seismogram_count(self, session: Session) -> None: Args: session: The database session. """ - create_snapshot(session) + active_event = get_active_event(session) + create_snapshot(session, active_event) snapshot = session.exec(select(AimbatSnapshot)).one() session.refresh(snapshot) assert snapshot.seismogram_count == len( @@ -265,7 +269,8 @@ def test_snapshot_selected_seismogram_count(self, session: Session) -> None: Args: session: The database session. """ - create_snapshot(session) + active_event = get_active_event(session) + create_snapshot(session, active_event) snapshot = session.exec(select(AimbatSnapshot)).one() session.refresh(snapshot) expected = sum(1 for s in snapshot.seismogram_parameters_snapshots if s.select) @@ -277,7 +282,8 @@ def test_snapshot_flipped_seismogram_count(self, session: Session) -> None: Args: session: The database session. """ - create_snapshot(session) + active_event = get_active_event(session) + create_snapshot(session, active_event) snapshot = session.exec(select(AimbatSnapshot)).one() session.refresh(snapshot) expected = sum(1 for s in snapshot.seismogram_parameters_snapshots if s.flip) @@ -308,7 +314,7 @@ def test_snapshot_counts_reflect_toggled_flip_and_select( session.add(to_deselect.parameters) session.commit() - create_snapshot(session) + create_snapshot(session, active_event) snapshot = session.exec(select(AimbatSnapshot)).one() session.refresh(snapshot) @@ -365,7 +371,7 @@ def test_snapshots_deleted(self, session: Session, event: AimbatEvent) -> None: session: The database session. event: An AimbatEvent to delete. """ - create_snapshot(session) + create_snapshot(session, event) session.refresh(event) assert len(event.snapshots) > 0 snapshot_ids = [s.id for s in event.snapshots] @@ -385,7 +391,7 @@ def test_snapshot_parameter_snapshots_deleted( session: The database session. event: An AimbatEvent to delete. """ - create_snapshot(session) + create_snapshot(session, event) session.refresh(event) session.delete(event) @@ -497,7 +503,8 @@ def test_parameter_snapshots_deleted( session: The database session. seismogram: An AimbatSeismogram to delete. """ - create_snapshot(session) + active_event = get_active_event(session) + create_snapshot(session, active_event) parameters_id = seismogram.parameters.id session.delete(seismogram) @@ -517,7 +524,8 @@ def test_event_parameters_snapshot_deleted(self, session: Session) -> None: Args: session: The database session. """ - create_snapshot(session) + active_event = get_active_event(session) + create_snapshot(session, active_event) snapshot = session.exec(select(AimbatSnapshot)).one() ep_snapshot_id = snapshot.event_parameters_snapshot.id @@ -532,7 +540,8 @@ def test_seismogram_parameters_snapshots_deleted(self, session: Session) -> None Args: session: The database session. """ - create_snapshot(session) + active_event = get_active_event(session) + create_snapshot(session, active_event) snapshot = session.exec(select(AimbatSnapshot)).one() sp_snapshot_ids = [s.id for s in snapshot.seismogram_parameters_snapshots] assert len(sp_snapshot_ids) > 0 diff --git a/uv.lock b/uv.lock index e5ceeb4d..cbdfc465 100644 --- a/uv.lock +++ b/uv.lock @@ -19,6 +19,7 @@ dependencies = [ { name = "matplotlib" }, { name = "pandas" }, { name = "pandas-stubs" }, + { name = "prompt-toolkit" }, { name = "pydantic-settings" }, { name = "pysmo" }, { name = "rich" }, @@ -55,6 +56,7 @@ requires-dist = [ { name = "matplotlib", specifier = ">=3.10.6" }, { name = "pandas", specifier = ">=3.0.1" }, { name = "pandas-stubs", specifier = ">=3.0.0.260204" }, + { name = "prompt-toolkit", specifier = ">=3.0.52" }, { name = "pydantic-settings", specifier = ">=2.10.1" }, { name = "pysmo", git = "https://github.com/pysmo/pysmo?rev=master" }, { name = "rich", specifier = ">=13.9.4" }, @@ -1676,6 +1678,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, ] +[[package]] +name = "prompt-toolkit" +version = "3.0.52" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wcwidth" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a1/96/06e01a7b38dce6fe1db213e061a4602dd6032a8a97ef6c1a862537732421/prompt_toolkit-3.0.52.tar.gz", hash = "sha256:28cde192929c8e7321de85de1ddbe736f1375148b02f2e17edd840042b1be855", size = 434198, upload-time = "2025-08-27T15:24:02.057Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/84/03/0d3ce49e2505ae70cf43bc5bb3033955d2fc9f932163e84dc0779cc47f48/prompt_toolkit-3.0.52-py3-none-any.whl", hash = "sha256:9aac639a3bbd33284347de5ad8d68ecc044b91a762dc39b7c21095fcd6a19955", size = 391431, upload-time = "2025-08-27T15:23:59.498Z" }, +] + [[package]] name = "propcache" version = "0.4.1" @@ -2491,6 +2505,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/33/e8/e40370e6d74ddba47f002a32919d91310d6074130fe4e17dabcafc15cbf1/watchdog-6.0.0-py3-none-win_ia64.whl", hash = "sha256:a1914259fa9e1454315171103c6a30961236f508b9b623eae470268bbcc6a22f", size = 79067, upload-time = "2024-11-01T14:07:11.845Z" }, ] +[[package]] +name = "wcwidth" +version = "0.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/35/a2/8e3becb46433538a38726c948d3399905a4c7cabd0df578ede5dc51f0ec2/wcwidth-0.6.0.tar.gz", hash = "sha256:cdc4e4262d6ef9a1a57e018384cbeb1208d8abbc64176027e2c2455c81313159", size = 159684, upload-time = "2026-02-06T19:19:40.919Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/68/5a/199c59e0a824a3db2b89c5d2dade7ab5f9624dbf6448dc291b46d5ec94d3/wcwidth-0.6.0-py3-none-any.whl", hash = "sha256:1a3a1e510b553315f8e146c54764f4fb6264ffad731b3d78088cdb1478ffbdad", size = 94189, upload-time = "2026-02-06T19:19:39.646Z" }, +] + [[package]] name = "win32-setctime" version = "1.2.0" diff --git a/zensical.toml b/zensical.toml index 3b3f9e74..b00cd2b3 100644 --- a/zensical.toml +++ b/zensical.toml @@ -18,9 +18,12 @@ nav = [ ] }, { "Usage" = [ {"Using AIMBAT" = "usage/index.md"}, - "usage/cli.md", - "usage/gui.md", - "usage/defaults.md", + {"CLI" = "usage/cli.md"}, + {"Shell" = "usage/shell.md"}, + {"TUI" = "usage/tui.md"}, + {"GUI" = "usage/gui.md"}, + {"API" = "usage/api.md"}, + {"Defaults" = "usage/defaults.md"}, ] }, { "API reference" = [ "api/aimbat.md",