From ec306cb9bed306cdbd636bd186cfb60103487059 Mon Sep 17 00:00:00 2001
From: Alexander Ji <alexji@uchicago.edu>
Date: Fri, 19 Apr 2024 22:00:13 +1000
Subject: [PATCH 1/2] Initialization. Nothing works.

---
 smh/radiative_transfer/korg/__init__.py  | 16 +++++
 smh/radiative_transfer/korg/cog.py       | 67 +++++++++++++++++++
 smh/radiative_transfer/korg/synthesis.py | 83 ++++++++++++++++++++++++
 smh/radiative_transfer/korg/utils.py     | 25 +++++++
 4 files changed, 191 insertions(+)
 create mode 100644 smh/radiative_transfer/korg/__init__.py
 create mode 100644 smh/radiative_transfer/korg/cog.py
 create mode 100644 smh/radiative_transfer/korg/synthesis.py
 create mode 100644 smh/radiative_transfer/korg/utils.py

diff --git a/smh/radiative_transfer/korg/__init__.py b/smh/radiative_transfer/korg/__init__.py
new file mode 100644
index 00000000..8c75794d
--- /dev/null
+++ b/smh/radiative_transfer/korg/__init__.py
@@ -0,0 +1,16 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+""" A standardized interface to the Korg radiative transfer code. """
+
+from __future__ import (division, print_function, absolute_import,
+                        unicode_literals)
+
+__all__ = ["abundance_cog", "synthesize", "RTError"]
+
+# See stackoverflow.com/questions/19913653/no-unicode-in-all-for-a-packages-init
+# __all__ = [_.encode("ascii") for _ in __all__]
+
+from .cog import abundance_cog
+from .synthesis import synthesize
+from .utils import RTError
diff --git a/smh/radiative_transfer/korg/cog.py b/smh/radiative_transfer/korg/cog.py
new file mode 100644
index 00000000..ef148a5e
--- /dev/null
+++ b/smh/radiative_transfer/korg/cog.py
@@ -0,0 +1,67 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+""" 
+Functionality to derive abundances from a transition's curve-of-growth using the
+Korg radiative transfer code.
+This doesn't actually work yet, because Korg assumes everything is on the linear part of the COG.
+"""
+
+from __future__ import (division, print_function, absolute_import,
+                        unicode_literals)
+
+from juliacall import Main as jl
+jl.seval("using Korg")
+Korg = jl.Korg
+
+import logging
+import numpy as np
+import re
+import yaml
+from pkg_resources import resource_stream
+
+from . import utils
+from .utils import RTError
+from smh.utils import element_to_species
+from smh import linelists
+
+logger = logging.getLogger(__name__)
+
+
+def abundance_cog(photosphere, transitions, full_output=False, verbose=False,
+    twd=None, **kwargs):
+    """
+    Calculate atomic line abundances by interpolating the measured 
+    equivalent width from the curve-of-growth. 
+    This doesn't actually work yet, because Korg assumes everything is on the linear part of the COG.
+    
+    # :param photosphere:
+    #     A formatted photosphere.
+
+    # :param transitions:
+    #     A list of atomic transitions with measured equivalent widths.
+
+    # :param verbose: [optional]
+    #     Specify verbose flags to MOOG. This is primarily used for debugging.
+
+    """
+    
+    # Create a temporary directory.
+    path = utils.twd_path(twd=twd,**kwargs)
+
+    # Write out the photosphere.
+    model_in, lines_in = path("model.in"), path("lines.in")
+    ## TODO
+    photosphere.write(model_in, format="marcs")
+    ## TODO
+    transitions.write(lines_in, format="moog")
+    
+    atm = Korg.read_photosphere(model_in)
+    MH = TODO
+    linelist = Korg.read_linelist(lines_in, format="moog")
+    # The first argument specifies the default [X/H] abundance, the second lets you specify the default [alpha/H] abundance, and the third lets you specify individual abundances.
+    A_X = Korg.format_A_X(MH) # the need for this should go away
+    measured_EWs = transition["equivalent_width"]
+    logeps = Korg.Fit.ews_to_abundances(atm, linelist, A_X, measured_EWs)
+    # return logeps
+    raise NotImplementedError
diff --git a/smh/radiative_transfer/korg/synthesis.py b/smh/radiative_transfer/korg/synthesis.py
new file mode 100644
index 00000000..ba4356df
--- /dev/null
+++ b/smh/radiative_transfer/korg/synthesis.py
@@ -0,0 +1,83 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+""" Synthesis functionality with the Korg radiative transfer code. """
+
+from __future__ import (division, print_function, absolute_import,
+                        unicode_literals)
+
+
+import logging
+import numpy as np
+import yaml
+from pkg_resources import resource_stream
+
+from . import utils
+
+logger = logging.getLogger(__name__)
+
+def synthesize(photosphere, transitions, abundances=None, isotopes=None,
+    verbose=False, twd=None, **kwargs):
+    """
+    Sythesize a stellar spectrum given the model photosphere and list of
+    transitions provided with Korg.
+    
+    :param photosphere:
+        A formatted photosphere.
+
+    :param transitions:
+        A list of atomic transitions with measured equivalent widths.
+
+    :param abundances: [optional]
+        The non-scaled-Solar abundances to use in the synthesis.
+
+    :param isotopes: [optional]
+        The isotopic fractions to use in the synthesis.
+
+    :param verbose: [optional]
+        Specify verbose flags to Korg. This is primarily used for debugging.
+
+    """
+
+    # Create a temporary directory and write out the photoshere and transitions.
+    path = utils.twd_path(twd=twd,**kwargs)
+    model_in, lines_in = path("model.in"), path("lines.in")
+    ## TODO: can't currently write a correct MARCS format
+    photosphere.write(model_in, format="marcs")
+    transitions.write(lines_in, format="moog")
+    wavemin, wavemax = transitions[0]-1.0,transitions[-1]+1.0
+    atm = Korg.read_model_atmosphere(model_in)
+    ## TODO: this doesn't support broadening parameters or dissociation energies
+    ## TODO: Korg will automatically apply isotopic ratios when it reads, need to update this?
+    linelist = Korg.read_linelist(lines_in, format="moog")
+    # read_linelist(filename; format="vald", isotopic_abundances=Korg.isotopic_abundances)
+
+    
+    # Abundances.
+    #A_X = Korg.format_A_X(-0.5, Dict("C" => -0.25))
+    #solution = synthesize(atm, linelist, A_X, 5000, 5100)
+    
+    #synthesize(atm, linelist, A_X, λ_start, λ_stop, [λ_step=0.01]; kwargs... )
+    #synthesize(atm, linelist, A_X, wavelength_ranges; kwargs... )
+    
+    #sol = Korg.synthesize(...)
+    # sol.flux
+    # sol.wavelengths
+    # sol.cntm
+    #flux: the output spectrum
+    #cntm: the continuum at each wavelength
+    #alpha: the linear absorption coefficient at each wavelength and atmospheric layer a Matrix of size (layers x wavelengths)
+    #number_densities: A dictionary mapping Species to vectors of number densities at each atmospheric layer
+    #electron_number_density: the electron number density at each atmospheric layer
+    #wavelengths: The vacuum wavelengths (in Å) over which the synthesis was performed. If air_wavelengths=true this will not be the same as the input wavelengths.
+    #subspectra: A vector of ranges which can be used to index into flux to extract the spectrum for each range provided in wavelength_ranges. If you use the standard λ_start, λ_stop, λ_step arguments, this will be a vector containing only one range.
+
+    spectra = []
+    all_abundances = []
+    for abundances in all_abundances:
+        out = Korg.synthesize(atm, linelist, abundances, wavemin, wavemax)
+        meta = {}
+        spectra.append((out.wavelengths, out.flux/out.cntm, meta))
+    
+    #return spectra
+    raise NotImplementedError
diff --git a/smh/radiative_transfer/korg/utils.py b/smh/radiative_transfer/korg/utils.py
new file mode 100644
index 00000000..7a69e231
--- /dev/null
+++ b/smh/radiative_transfer/korg/utils.py
@@ -0,0 +1,25 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+""" Utility functions related to MOOG. """
+
+from __future__ import (division, print_function, absolute_import,
+                        unicode_literals)
+
+import logging
+import numpy as np
+import os
+import signal
+import subprocess
+#import tempfile
+
+from smh.photospheres.abundances import asplund_2009 as solar_composition
+from smh.utils import elems_isotopes_ion_to_species, element_to_atomic_number
+from smh.utils import mkdtemp
+from six import iteritems, string_types
+
+logger = logging.getLogger(__name__)
+
+class RTError(BaseException):
+    pass
+

From 685c7daf83bfbe562d8e59ec5833bbad2c6c591c Mon Sep 17 00:00:00 2001
From: Alexander Ji <alexji@uchicago.edu>
Date: Sun, 16 Jun 2024 13:29:49 -0500
Subject: [PATCH 2/2] Adding some stuff that I won't touch later

---
 smh/radiative_transfer/korg/cog.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/smh/radiative_transfer/korg/cog.py b/smh/radiative_transfer/korg/cog.py
index ef148a5e..a514764e 100644
--- a/smh/radiative_transfer/korg/cog.py
+++ b/smh/radiative_transfer/korg/cog.py
@@ -60,7 +60,7 @@ def abundance_cog(photosphere, transitions, full_output=False, verbose=False,
     MH = TODO
     linelist = Korg.read_linelist(lines_in, format="moog")
     # The first argument specifies the default [X/H] abundance, the second lets you specify the default [alpha/H] abundance, and the third lets you specify individual abundances.
-    A_X = Korg.format_A_X(MH) # the need for this should go away
+    A_X = Korg.format_A_X(MH) # the need for this should go away after including the new code, which shouldn't need an initialization?
     measured_EWs = transition["equivalent_width"]
     logeps = Korg.Fit.ews_to_abundances(atm, linelist, A_X, measured_EWs)
     # return logeps