Skip to content

Commit

Permalink
Added script arguments.
Browse files Browse the repository at this point in the history
  • Loading branch information
alxmrs committed Mar 23, 2024
1 parent e6bfffe commit f9c57ba
Show file tree
Hide file tree
Showing 3 changed files with 66 additions and 50 deletions.
50 changes: 0 additions & 50 deletions demo/sst

This file was deleted.

63 changes: 63 additions & 0 deletions demo/sst.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
#!/usr/bin/env python3
"""Demo of calculating global average sea surface temperature (SST) with SQL.
Please run the following to set up cloud resources:
```
gcloud auth application-default login
coiled setup
```
"""
import argparse
import xarray as xr
import xarray_sql as qr

parser = argparse.ArgumentParser()
parser.add_argument('--start', type=str, default='2020-01-01', help='start time ISO string')
parser.add_argument('--end', type=str, default='2020-01-02', help='end time ISO string')
parser.add_argument('--cluster', action='store_true', help='deploy on coiled cluster')

args = parser.parse_args()

if args.cluster:
from coiled import Cluster

cluster = Cluster(
region='us-central1',
worker_memory='16 GiB',
spot_policy='spot_with_fallback',
arm=True,
)
client = cluster.get_client()
cluster.adapt(minimum=1, maximum=100)
else:
from dask.distributed import LocalCluster
cluster = LocalCluster(processes=False)
client = cluster.get_client()

era5_ds = xr.open_zarr(
'gs://gcp-public-data-arco-era5/ar/'
'1959-2022-full_37-1h-0p25deg-chunk-1.zarr-v2',
chunks={'time': 240, 'level': 1}
)
print('dataset opened.')
era5_sst_ds = era5_ds[['sea_surface_temperature']].sel(
time=slice(args.start, args.end),
level=1000, # surface level only.
)

c = qr.Context()
# chunk sizes determined from VM memory limit of 16 GiB.
c.create_table('era5', era5_sst_ds, chunks=dict(time=24))

print('beginning query.')
df = c.sql("""
SELECT
DATE("time") as date,
AVG("sea_surface_temperature") as daily_avg_sst
FROM
"era5"
GROUP BY
DATE("time")
""")

df.to_csv(f'global_avg_sst_{args.start}-{args.end}_*.cvs')
3 changes: 3 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,9 @@ dev = [
"pyink",
"py-spy"
]
demo = [
"coiled"
]

[project.urls]
Homepage = "https://github.com/alxmrs/xarray-sql"
Expand Down

0 comments on commit f9c57ba

Please sign in to comment.