diff --git a/data/trexio.json b/data/trexio.json deleted file mode 100644 index 9bdc263..0000000 --- a/data/trexio.json +++ /dev/null @@ -1,231 +0,0 @@ -{ - - "metadata": { - "code_num" : [ "dim", [] ] - , "code" : [ "str", [ "metadata.code_num" ] ] - , "author_num" : [ "dim", [] ] - , "author" : [ "str", [ "metadata.author_num" ] ] - , "package_version" : [ "str", [] ] - , "description" : [ "str", [] ] - , "unsafe" : [ "int", [] ] - } , - - "nucleus": { - "num" : [ "dim" , [] ] - , "charge" : [ "float", [ "nucleus.num" ] ] - , "coord" : [ "float", [ "nucleus.num", "3" ] ] - , "label" : [ "str" , [ "nucleus.num" ] ] - , "point_group" : [ "str" , [] ] - , "repulsion" : [ "float", [] ] - } , - - "cell": { - "a" : [ "float", [ "3" ] ] - , "b" : [ "float", [ "3" ] ] - , "c" : [ "float", [ "3" ] ] - , "G_a" : [ "float", [ "3" ] ] - , "G_b" : [ "float", [ "3" ] ] - , "G_c" : [ "float", [ "3" ] ] - , "two_pi" : [ "int" , [] ] - } , - - "pbc": { - "periodic" : [ "int" , [] ] - , "k_point" : [ "float", [ "3" ] ] - } , - - "electron": { - "num" : [ "dim", [] ] - , "up_num" : [ "int", [] ] - , "dn_num" : [ "int", [] ] - } , - - "state": { - "num" : [ "dim" , [] ] - , "id" : [ "index", [] ] - , "energy" : [ "float", [] ] - , "current_label" : [ "str" , [] ] - , "label" : [ "str" , [ "state.num" ] ] - , "file_name" : [ "str" , [ "state.num" ] ] - } , - - "basis": { - "type" : [ "str" , [] ] - , "prim_num" : [ "dim" , [] ] - , "shell_num" : [ "dim" , [] ] - , "nao_grid_num" : [ "dim" , [] ] - , "interp_coeff_cnt" : [ "dim" , [] ] - , "nucleus_index" : [ "index", [ "basis.shell_num" ] ] - , "shell_ang_mom" : [ "int" , [ "basis.shell_num" ] ] - , "shell_factor" : [ "float", [ "basis.shell_num" ] ] - , "r_power" : [ "int" , [ "basis.shell_num" ] ] - , "nao_grid_start" : [ "index", [ "basis.shell_num" ] ] - , "nao_grid_size" : [ "dim" , [ "basis.shell_num" ] ] - , "shell_index" : [ "index", [ "basis.prim_num" ] ] - , "exponent" : [ "float", [ "basis.prim_num" ] ] - , "coefficient" : [ "float", [ "basis.prim_num" ] ] - , "prim_factor" : [ "float", [ "basis.prim_num" ] ] - , "e_cut" : [ "float", [] ] - , "nao_grid_radius" : [ "float", [ "basis.nao_grid_num" ] ] - , "nao_grid_phi" : [ "float", [ "basis.nao_grid_num" ] ] - , "nao_grid_grad" : [ "float", [ "basis.nao_grid_num" ] ] - , "nao_grid_lap" : [ "float", [ "basis.nao_grid_num" ] ] - , "interpolator_kind" : [ "str" , [] ] - , "interpolator_phi" : [ "float", [ "basis.nao_grid_num", "basis.interp_coeff_cnt" ] ] - , "interpolator_grad" : [ "float", [ "basis.nao_grid_num", "basis.interp_coeff_cnt" ] ] - , "interpolator_lap" : [ "float", [ "basis.nao_grid_num", "basis.interp_coeff_cnt" ] ] - } , - - "ecp": { - "max_ang_mom_plus_1" : [ "int" , [ "nucleus.num" ] ] - , "z_core" : [ "int" , [ "nucleus.num" ] ] - , "num" : [ "dim" , [] ] - , "ang_mom" : [ "int" , [ "ecp.num" ] ] - , "nucleus_index" : [ "index", [ "ecp.num" ] ] - , "exponent" : [ "float", [ "ecp.num" ] ] - , "coefficient" : [ "float", [ "ecp.num" ] ] - , "power" : [ "int" , [ "ecp.num" ] ] - } , - - "grid": { - "description" : [ "str" , [] ] - , "rad_precision" : [ "float", [] ] - , "num" : [ "dim" , [] ] - , "max_ang_num" : [ "int" , [] ] - , "min_ang_num" : [ "int" , [] ] - , "coord" : [ "float", [ "grid.num" ] ] - , "weight" : [ "float", [ "grid.num" ] ] - , "ang_num" : [ "dim" , [] ] - , "ang_coord" : [ "float", [ "grid.ang_num" ] ] - , "ang_weight" : [ "float", [ "grid.ang_num" ] ] - , "rad_num" : [ "dim" , [] ] - , "rad_coord" : [ "float", [ "grid.rad_num" ] ] - , "rad_weight" : [ "float", [ "grid.rad_num" ] ] - } , - - "ao": { - "cartesian" : [ "int" , [] ] - , "num" : [ "dim" , [] ] - , "shell" : [ "index", [ "ao.num" ] ] - , "normalization" : [ "float", [ "ao.num" ] ] - } , - - "ao_1e_int": { - "overlap" : [ "float", [ "ao.num", "ao.num" ] ] - , "kinetic" : [ "float", [ "ao.num", "ao.num" ] ] - , "potential_n_e" : [ "float", [ "ao.num", "ao.num" ] ] - , "ecp" : [ "float", [ "ao.num", "ao.num" ] ] - , "core_hamiltonian" : [ "float", [ "ao.num", "ao.num" ] ] - , "overlap_im" : [ "float", [ "ao.num", "ao.num" ] ] - , "kinetic_im" : [ "float", [ "ao.num", "ao.num" ] ] - , "potential_n_e_im" : [ "float", [ "ao.num", "ao.num" ] ] - , "ecp_im" : [ "float", [ "ao.num", "ao.num" ] ] - , "core_hamiltonian_im" : [ "float", [ "ao.num", "ao.num" ] ] - } , - - "ao_2e_int": { - "eri" : [ "float sparse", [ "ao.num", "ao.num", "ao.num", "ao.num" ] ] - , "eri_lr" : [ "float sparse", [ "ao.num", "ao.num", "ao.num", "ao.num" ] ] - , "eri_cholesky_num" : [ "dim" , [] ] - , "eri_cholesky" : [ "float sparse", [ "ao_2e_int.eri_cholesky_num", "ao.num", "ao.num" ] ] - , "eri_lr_cholesky_num" : [ "dim" , [] ] - , "eri_lr_cholesky" : [ "float sparse", [ "ao_2e_int.eri_lr_cholesky_num", "ao.num", "ao.num" ] ] - } , - - "mo": { - "type" : [ "str" , [] ] - , "num" : [ "dim" , [] ] - , "coefficient" : [ "float", [ "mo.num", "ao.num" ] ] - , "coefficient_im" : [ "float", [ "mo.num", "ao.num" ] ] - , "class" : [ "str" , [ "mo.num" ] ] - , "symmetry" : [ "str" , [ "mo.num" ] ] - , "occupation" : [ "float", [ "mo.num" ] ] - , "energy" : [ "float", [ "mo.num" ] ] - , "spin" : [ "int" , [ "mo.num" ] ] - } , - - "mo_1e_int": { - "overlap" : [ "float", [ "mo.num", "mo.num" ] ] - , "kinetic" : [ "float", [ "mo.num", "mo.num" ] ] - , "potential_n_e" : [ "float", [ "mo.num", "mo.num" ] ] - , "ecp" : [ "float", [ "mo.num", "mo.num" ] ] - , "core_hamiltonian" : [ "float", [ "mo.num", "mo.num" ] ] - , "overlap_im" : [ "float", [ "mo.num", "mo.num" ] ] - , "kinetic_im" : [ "float", [ "mo.num", "mo.num" ] ] - , "potential_n_e_im" : [ "float", [ "mo.num", "mo.num" ] ] - , "ecp_im" : [ "float", [ "mo.num", "mo.num" ] ] - , "core_hamiltonian_im" : [ "float", [ "mo.num", "mo.num" ] ] - } , - - "mo_2e_int": { - "eri" : [ "float sparse", [ "mo.num", "mo.num", "mo.num", "mo.num" ] ] - , "eri_lr" : [ "float sparse", [ "mo.num", "mo.num", "mo.num", "mo.num" ] ] - , "eri_cholesky_num" : [ "dim" , [] ] - , "eri_cholesky" : [ "float sparse", [ "mo_2e_int.eri_cholesky_num", "mo.num", "mo.num" ] ] - , "eri_lr_cholesky_num" : [ "dim" , [] ] - , "eri_lr_cholesky" : [ "float sparse", [ "mo_2e_int.eri_lr_cholesky_num", "mo.num", "mo.num" ] ] - } , - - "determinant": { - "num" : [ "dim readonly" , [] ] - , "list" : [ "int special" , [ "determinant.num" ] ] - , "coefficient" : [ "float buffered", [ "determinant.num" ] ] - } , - - "csf": { - "num" : [ "dim readonly" , [] ] - , "coefficient" : [ "float buffered", [ "csf.num" ] ] - , "det_coefficient" : [ "float sparse" , [ "csf.num", "determinant.num" ] ] - } , - - "amplitude": { - "single" : [ "float sparse", [ "mo.num", "mo.num" ] ] - , "single_exp" : [ "float sparse", [ "mo.num", "mo.num" ] ] - , "double" : [ "float sparse", [ "mo.num", "mo.num", "mo.num", "mo.num" ] ] - , "double_exp" : [ "float sparse", [ "mo.num", "mo.num", "mo.num", "mo.num" ] ] - , "triple" : [ "float sparse", [ "mo.num", "mo.num", "mo.num", "mo.num", "mo.num", "mo.num" ] ] - , "triple_exp" : [ "float sparse", [ "mo.num", "mo.num", "mo.num", "mo.num", "mo.num", "mo.num" ] ] - , "quadruple" : [ "float sparse", [ "mo.num", "mo.num", "mo.num", "mo.num", "mo.num", "mo.num", "mo.num", "mo.num" ] ] - , "quadruple_exp" : [ "float sparse", [ "mo.num", "mo.num", "mo.num", "mo.num", "mo.num", "mo.num", "mo.num", "mo.num" ] ] - } , - - "rdm": { - "1e" : [ "float" , [ "mo.num", "mo.num" ] ] - , "1e_up" : [ "float" , [ "mo.num", "mo.num" ] ] - , "1e_dn" : [ "float" , [ "mo.num", "mo.num" ] ] - , "2e" : [ "float sparse", [ "mo.num", "mo.num", "mo.num", "mo.num" ] ] - , "2e_upup" : [ "float sparse", [ "mo.num", "mo.num", "mo.num", "mo.num" ] ] - , "2e_dndn" : [ "float sparse", [ "mo.num", "mo.num", "mo.num", "mo.num" ] ] - , "2e_updn" : [ "float sparse", [ "mo.num", "mo.num", "mo.num", "mo.num" ] ] - , "2e_cholesky_num" : [ "dim" , [] ] - , "2e_cholesky" : [ "float sparse", [ "rdm.2e_cholesky_num", "mo.num", "mo.num" ] ] - , "2e_upup_cholesky_num" : [ "dim" , [] ] - , "2e_upup_cholesky" : [ "float sparse", [ "rdm.2e_upup_cholesky_num", "mo.num", "mo.num" ] ] - , "2e_dndn_cholesky_num" : [ "dim" , [] ] - , "2e_dndn_cholesky" : [ "float sparse", [ "rdm.2e_dndn_cholesky_num", "mo.num", "mo.num" ] ] - , "2e_updn_cholesky_num" : [ "dim" , [] ] - , "2e_updn_cholesky" : [ "float sparse", [ "rdm.2e_updn_cholesky_num", "mo.num", "mo.num" ] ] - } , - - "jastrow": { - "type" : [ "str" , [] ] - , "en_num" : [ "dim" , [] ] - , "ee_num" : [ "dim" , [] ] - , "een_num" : [ "dim" , [] ] - , "en" : [ "float" , [ "jastrow.en_num" ] ] - , "ee" : [ "float" , [ "jastrow.ee_num" ] ] - , "een" : [ "float" , [ "jastrow.een_num" ] ] - , "en_nucleus" : [ "index" , [ "jastrow.en_num" ] ] - , "een_nucleus" : [ "index" , [ "jastrow.een_num" ] ] - , "ee_scaling" : [ "float" , [] ] - , "en_scaling" : [ "float" , [ "nucleus.num" ] ] - } , - - "qmc": { - "num" : [ "dim" , [] ] - , "point" : [ "float", [ "qmc.num", "electron.num", "3" ] ] - , "psi" : [ "float", [ "qmc.num" ] ] - , "e_loc" : [ "float", [ "qmc.num" ] ] - } - - } \ No newline at end of file diff --git a/flake.lock b/flake.lock index baa34c4..d334c34 100644 --- a/flake.lock +++ b/flake.lock @@ -38,11 +38,11 @@ }, "nixpkgs": { "locked": { - "lastModified": 1733932923, - "narHash": "sha256-6rB4DwsPBntgP9PhObgZtBDWlNBJ/hpHjwuBUUHJZJI=", + "lastModified": 1734359469, + "narHash": "sha256-NuYABf+URWfc0LHSv5nYNhbh0hW+E+U4vjxnFuv7K2I=", "owner": "NixOS", "repo": "nixpkgs", - "rev": "f5c423099d9d1702aeee68e5b94f6a9836890274", + "rev": "3175ce85131b128562f362780ee538e37dde2fab", "type": "github" }, "original": { diff --git a/get-json-spec.sh b/get-json-spec.sh new file mode 100755 index 0000000..466a8bd --- /dev/null +++ b/get-json-spec.sh @@ -0,0 +1,13 @@ +#! /usr/bin/env bash + +# Create a temporary C file, that merely imports trexio.h +TREXIO_TMP=$(mktemp --suffix=.c) +trap "rm -f $TREXIO_TMP" EXIT ERR + +# Run the C preprocessor on that file to get the path to trexio.h +echo "#include " > $TREXIO_TMP +TREXIO_HEADER_PATH=$(gcc -E test.c | grep '^# [0-9]\+ "' | grep 'trexio\.h' | head -n 1 | awk -F'"' '{print $2}') + +# Extract the JSON specification from a comment in the header file +sed -n '/\/\* JSON configuration/,/\*\//p' $TREXIO_HEADER_PATH | sed '1d;$d' + diff --git a/nix/overlay.nix b/nix/overlay.nix index cf55780..9f6f308 100644 --- a/nix/overlay.nix +++ b/nix/overlay.nix @@ -8,7 +8,9 @@ final: prev: { in prev.haskell.packageOverrides hfinal hprev // { - trexio-hs = hfinal.callCabal2nix "trexio-hs" ../. { inherit (final) trexio; }; + trexio-hs = hfinal.callCabal2nix "trexio-hs" ../. { + inherit (final) trexio; + }; }; }; } diff --git a/src/TREXIO/Internal/TH.hs b/src/TREXIO/Internal/TH.hs index bbae2fa..66c2a24 100644 --- a/src/TREXIO/Internal/TH.hs +++ b/src/TREXIO/Internal/TH.hs @@ -15,7 +15,7 @@ import Data.Char import Data.Coerce import Data.Map (Map) import Data.Map qualified as Map -import Data.Massiv.Array as Massiv hiding (Dim, forM, forM_, mapM, product, replicate, toList, zip, throwM) +import Data.Massiv.Array as Massiv hiding (Dim, forM, forM_, mapM, product, replicate, throwM, toList, zip) import Data.Massiv.Array qualified as Massiv import Data.Massiv.Array.Manifest.Vector qualified as Massiv import Data.Massiv.Array.Unsafe (unsafeWithPtr) @@ -23,13 +23,14 @@ import Data.Maybe import Data.Text (Text) import Data.Text qualified as T import Data.Vector qualified as V -import Foreign hiding (peekArray, withArray) +import Foreign hiding (peekArray, void, withArray) import Foreign.C.ConstPtr import Foreign.C.String import Foreign.C.Types import GHC.Generics (Generic) import Language.Haskell.TH import Language.Haskell.TH.Syntax (Lift (..)) +import System.Process.Typed import TREXIO.CooArray import TREXIO.Internal.Base import TREXIO.Internal.Marshaller @@ -41,6 +42,21 @@ tshow = T.pack . show -------------------------------------------------------------------------------- +{- | Attempts to obtain the JSON specification from the trexio.h header. Magic +happens in a bash script +-} +getJsonSpec :: (MonadIO m, MonadThrow m) => m TrexioScheme +getJsonSpec = do + (ec, stdout, stderr) <- readProcess . shell $ "./get-json-spec.sh" + jsonSpec <- case ec of + ExitSuccess -> return stdout + ExitFailure _ -> throwString . show $ stderr + case eitherDecode jsonSpec of + Left err -> throwString err + Right spec -> return spec + +-------------------------------------------------------------------------------- + {- | The overall data structure TREXIO uses to represent a wave function as a JSON specification. A TREXIO scheme consists of multiple data groups and each data group has multiple fields. A field may require knowledge of other fields. @@ -659,6 +675,20 @@ mkReadFns groupName dataName fieldType = case dims of |] | otherwise -> error $ "mkReadFns: unsupported field type for 3D data: " <> show fieldType [d1, d2, d3, d4] + | isFloatField fieldType -> + [e| + \trexio -> + liftIO $ do + sz1 <- $(mkSizeFn d1) trexio + sz2 <- $(mkSizeFn d2) trexio + sz3 <- $(mkSizeFn d3) trexio + sz4 <- $(mkSizeFn d4) trexio + allocaArray (sz1 * sz2 * sz3 * sz4) $ \buf -> do + ec <- exitCodeH <$> $(varE . mkName $ mkCFnName Read groupName dataName) trexio buf + case ec of + Success -> peekArray (Sz4 sz1 sz2 sz3 sz4) (castPtr buf) + _ -> throwM ec + |] | isSparseFloat fieldType -> [e| \trexio -> liftIO $ do @@ -982,6 +1012,16 @@ mkWriteFns scheme groupName dataName fieldType = case dims of |] | otherwise -> error $ "mkWriteFns: unsupported field type for 3D data: " <> show fieldType [d1, d2, d3, d4] + | isFloatField fieldType -> + [e| + \trexio arr -> liftIO . unsafeWithPtr arr $ \arrPtr -> do + let Sz4 sz1 sz2 sz3 sz4 = size arr + $(mkWriteSzFn scheme d1) trexio sz1 + $(mkWriteSzFn scheme d2) trexio sz2 + $(mkWriteSzFn scheme d3) trexio sz3 + $(mkWriteSzFn scheme d4) trexio sz4 + checkEC $ $(varE . mkName $ mkCFnName Write groupName dataName) trexio (castPtr arrPtr) + |] | isSparseFloat fieldType -> [e| \trexio cooArr -> liftIO $ do diff --git a/src/TREXIO/LowLevel/Scheme.hs b/src/TREXIO/LowLevel/Scheme.hs index 457622c..49eb0e6 100644 --- a/src/TREXIO/LowLevel/Scheme.hs +++ b/src/TREXIO/LowLevel/Scheme.hs @@ -19,6 +19,6 @@ import Language.Haskell.TH.Syntax (lift) scheme :: TrexioScheme scheme = $(do - Just trexio <- runIO $ decodeFileStrict @TrexioScheme "./data/trexio.json" + trexio <- runIO getJsonSpec lift trexio ) \ No newline at end of file diff --git a/trexio-hs.cabal b/trexio-hs.cabal index 289a6ef..6a48feb 100644 --- a/trexio-hs.cabal +++ b/trexio-hs.cabal @@ -16,7 +16,7 @@ maintainer: phillip.seeber@uni-jena.de category: Data build-type: Simple extra-doc-files: CHANGELOG.md -data-files: data/trexio.json +data-files: get-json-spec.sh tested-with: GHC == {9.6.6, 9.8.2} description: This package provides low- and high-level Haskell bindings for [TREXIO, a portable file format for storing wave function data](https://trex-coe.github.io/trexio/). @@ -81,6 +81,7 @@ common deps massiv >= 1.0.0.0 && < 1.1, safe-exceptions >= 0.1.7 && < 0.2, template-haskell >= 2.20 && < 2.23, + typed-process >= 0.2.12 && < 0.3, text >= 2.0 && < 2.2, vector >= 0.13 && < 0.14 @@ -147,6 +148,6 @@ test-suite trexio-test tasty >= 1.4 && < 1.6, tasty-hunit >= 0.10 && < 0.11, tasty-hedgehog >= 1.4 && < 1.5, - hedgehog >= 1.4 && < 1.6, - temporary >= 1.3 && < 1.4 + temporary >= 1.3 && < 1.4, + hedgehog >= 1.4 && < 1.6