Automate installation of SOCRATES (#48)

* Load socrates env in Python

* Set default directory

* Add tool for downloading SOCRATES

* Move set_rad_env away from __init__

* Clean download code

* Turn set_rad_env into a singleton

* Revise dowload function to download source and compile

* Fix GA

* Add platformdirs dependency

* Set path to socrates dir

* Fix path to SOCRATES

* Add documentation

* Format code

* Update documentation

.github/workflows/tests.yaml

@@ -22,7 +22,8 @@ jobs:
 
       - name: Setup system
         run: |
-          sudo apt-get install libnetcdff-dev netcdf-bin gfortran gcc
+          sudo apt-get update
+          sudo apt-get install libnetcdff-dev netcdf-bin
 
       - name: Get SOCRATES
         uses: actions/checkout@v4
@@ -45,7 +46,7 @@ jobs:
           export LD_LIBRARY_PATH=""
           cd SOCRATES
           ./configure
-          ./build_code 
+          ./build_code
           cd ..
 
       - name: Set up Python ${{ matrix.python-version }}
@@ -73,7 +74,7 @@ jobs:
       - name: Test with pytest
         run: |
           export FWL_DATA="/home/runner/work/fwl_data"
-          source SOCRATES/set_rad_env
+          export SOCRATES="/home/runner/work/JANUS/JANUS/SOCRATES"
           coverage run -m pytest
 
       - name: Report coverage
@@ -98,5 +99,3 @@ jobs:
           minColorRange: 50
           maxColorRange: 90
           valColorRange: ${{ env.total }}
-
-

docs/getting_started.md

@@ -1,8 +1,39 @@
-This page shows you how to get started using JANUS.
-
-Lorem ipsum dolor sit amet, consectetur adipisicing elit, sed do eiusmod
-tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam,
-quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo
-consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse
-cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non
-proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
+# Getting started
+
+JANUS contains a small CLI tool to help get set up with JANUS.
+
+## Install SOCRATES
+
+Download and install SOCRATES:
+
+```console
+janus download socrates
+```
+
+Make sure you have the netcdf fortran libraries installed:
+
+```
+sudo apt install libnetcdff-dev netcdf-bin
+```
+
+## Download data
+
+Download spectral and stellar data:
+
+```console
+janus download spectral
+janus download stellar
+```
+
+## Environment variables
+
+### `SOCRATES`
+
+By default, SOCRATES is installed to the default location based on the [XDG specification](https://specifications.freedesktop.org/basedir-spec/latest/).
+
+If you install and compile [SOCRATES](https://github.com/nichollsh/SOCRATES) yourself,
+you can override the path using the `SOCRATES` environment variable, e.g.
+
+```console
+SOCRATES=/home/user/path/to/SOCRATES pytest
+```

pyproject.toml

@@ -34,6 +34,8 @@ dependencies = [
   'numpy',
   'osfclient',
   'pandas',
+  'platformdirs',
+  'requests',
   'scipy',
   'seaborn',
   'toml',

src/janus/cli.py

@@ -1,14 +1,17 @@
 import click
 
+
 @click.group()
 def cli():
     pass
 
+
 @click.group()
 def download():
     """Download data and dependencies"""
     pass
 
+
 @click.command()
 @click.option('-n', '--name', 'fname', type=str, help='Name of the spectra')
 @click.option('-b', '--band', 'nband', type=int, help='Number of the band', default=256)
@@ -18,18 +21,25 @@ def spectral(**kwargs):
     By default, download all files.
     """
     from .utils.data import DownloadSpectralFiles
+
     DownloadSpectralFiles(**kwargs)
 
+
 @click.command()
 def stellar():
     """Download stellar spectra"""
     from .utils.data import DownloadStellarSpectra
+
     DownloadStellarSpectra()
 
+
 @click.command()
-def socrates():
+@click.option('-r', '--ref', type=str, default='main', help='Git reference (hash) to download')
+def socrates(**kwargs):
     """Download SOCRATES code"""
-    raise NotImplementedError
+    from .socrates import download_socrates
+
+    download_socrates(**kwargs)
 
 
 cli.add_command(download)
@@ -38,4 +48,4 @@ def socrates():
 download.add_command(socrates)
 
 if __name__ == '__main__':
-    cli()
+    cli()

src/janus/set_socrates_env.py

@@ -0,0 +1,41 @@
+"""Import this file to set environment variables for socrates.
+
+Based on:
+https://github.com/FormingWorlds/SOCRATES/blob/main/sbin/set_rad_env_tmp
+"""
+
+from __future__ import annotations
+
+from .socrates import SOCRATES_DIR
+
+import os
+import sys
+import zipfile
+from pathlib import Path
+
+import click
+import platformdirs
+import requests
+
+if not SOCRATES_DIR.exists():
+    raise RuntimeError(f'Cannot find SOCRATES in this location: {SOCRATES_DIR}')
+
+with open(SOCRATES_DIR / 'version') as f:
+    version = f.readline()
+
+print(f'socrates location: {SOCRATES_DIR}')
+print(f'socrates version: {version}')
+
+sep = os.pathsep
+
+os.environ['RAD_DIR'] = str(SOCRATES_DIR)
+os.environ['RAD_BIN'] = str(SOCRATES_DIR / 'bin')
+os.environ['RAD_DATA'] = str(SOCRATES_DIR / 'data')
+os.environ['RAD_SCRIPT'] = str(SOCRATES_DIR / 'sbin')
+os.environ['LOCK_FILE'] = 'radiation_code.lock'
+os.environ['PATH'] = str(SOCRATES_DIR / 'bin') + sep + os.environ['PATH']
+os.environ['PATH'] = str(SOCRATES_DIR / 'sbin') + sep + os.environ['PATH']
+os.environ['MANPATH'] = str(SOCRATES_DIR / 'man') + sep + os.environ.get('MANPATH', '')
+sys.path.append(str(SOCRATES_DIR / 'python'))
+
+os.environ['LD_LIBRARY_PATH'] = 'netcdfff' + sep + os.environ.get('LD_LIBRARY_PATH', '')

src/janus/socrates.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+import os
+import subprocess as sp
+import zipfile
+from pathlib import Path
+
+import click
+import platformdirs
+import requests
+
+SOCRATES_DATA_DIR = Path(platformdirs.user_data_dir('socrates'))
+SOCRATES_DIR = Path(os.environ.get('SOCRATES', SOCRATES_DATA_DIR / 'SOCRATES'))
+
+
+def _set_permissions(drc: Path):
+    """Set executable flag for scripts."""
+    # Set executable permissions for make script
+    (drc / 'make' / 'mkdep').chmod(mode=33261)
+
+
+def _download(*, url: str, filename: str):
+    """Download file from url."""
+    chunk_size = 1024 * 8
+
+    with requests.get(url, stream=True) as r:
+        r.raise_for_status()
+        total_size = int(r.headers.get('Content-Length', 0))
+        with (
+            open(filename, 'wb') as f,
+            click.progressbar(label='Downloading', length=total_size) as pbar,
+        ):
+            for chunk in r.iter_content(chunk_size=chunk_size):
+                f.write(chunk)
+                pbar.update(chunk_size)
+
+
+def download_socrates(ref: str = 'main'):
+    """Version can be 'main' or the git commit hash."""
+    filename = 'socrates.zip'
+    path = 'refs/heads/main' if ref == 'main' else ref
+
+    _download(
+        url=f'https://github.com/nichollsh/SOCRATES/archive/{path}.zip',
+        filename=filename,
+    )
+
+    SOCRATES_DATA_DIR.mkdir(exist_ok=True, parents=True)
+
+    subdir = f'SOCRATES-{ref}'
+    target_dir = SOCRATES_DATA_DIR / subdir
+
+    click.echo(f'Extracting to {target_dir}')
+
+    with zipfile.ZipFile(filename, 'r') as zip_ref:
+        zip_ref.extractall(SOCRATES_DATA_DIR)
+
+    _set_permissions(drc=target_dir)
+
+    sp.run(['bash', 'configure'], cwd=target_dir)
+    sp.run(['bash', 'build_code'], cwd=target_dir)
+
+    symlink = SOCRATES_DATA_DIR / 'SOCRATES'
+    symlink.unlink(missing_ok=True)
+    symlink.symlink_to(target_dir)
+
+    print(f'SOCRATES downloaded to {target_dir}')

src/janus/utils/StellarSpectrum.py

@@ -11,6 +11,9 @@
 import subprocess
 from scipy.interpolate import PchipInterpolator
 
+from .. import set_socrates_env  # noqa
+
+
 def PrepareStellarSpectrum(wl, fl, star_file, nbins_max=95000):
     """Write a stellar spectrum.
 

src/janus/utils/socrates.py

@@ -17,6 +17,8 @@
 from janus.utils.atmosphere_column import atmos
 import janus.utils.phys as phys
 
+from .. import set_socrates_env  # noqa
+
 
 def radCompSoc(atm, dirs, recalc, rscatter=False,
                rewrite_cfg=True, rewrite_tmp=True, rewrite_gas=False):

tests/conftest.py

@@ -1,9 +1,5 @@
 import os
 
-if not os.environ.get('RAD_DIR'):
-    raise Exception(
-        'Socrates environment variables not set! Have you installed Socrates and sourced set_rad_env?'
-    )
 if not os.environ.get('FWL_DATA'):
     raise Exception(
         'The FWL_DATA environment variable where spectral and evolution tracks data will be downloaded needs to be set up!'