diff --git a/docs/make.jl b/docs/make.jl index 4eb2db8..f0a040d 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -29,8 +29,17 @@ makedocs(; "Fitting an Electronic Friction Tensor" => "fitting-eft.md", "Fitting a Dissipative Particle Dynamics Friction Model" => "fitting-mbdpd.md" ], - "Function Manual" => Any[ - "function-manual.md", + # "Importing & Exporting Friction Data" => Any[ + # "data-import-export.md", + # ], + # "Function Manual" => Any[ + # "function-manual.md" + # ] + "Function Manual" => Any[ + "ACEfriction.FrictionModels" => "./function-manual/ACEfriction.FrictionModels.md", + "ACEfriction.MatrixModels" => "./function-manual/ACEfriction.MatrixModels.md", + "ACEfriction.FrictionFit" => "./function-manual/ACEfriction.FrictionFit.md", + "ACEfriction.DataUtils" => "./function-manual/ACEfriction.DataUtils.md" ] ] ) diff --git a/docs/runLiveServer.jl b/docs/runLiveServer.jl new file mode 100644 index 0000000..e69de29 diff --git a/docs/src/data-import-export.md b/docs/src/data-import-export.md new file mode 100644 index 0000000..70693e9 --- /dev/null +++ b/docs/src/data-import-export.md @@ -0,0 +1,71 @@ +# Friction Data + +To train a friction tensor model in ACEfriciton.jl, friction data +$D = (R_i,\Gamma_i)_{i=1}^{N_{\rm obs}}$ comprised of atomic configurations, $R_i$, and a friction tensors, $\Gamma_i$ is required. + +`ACEfriction.jl` implements the structure `FrictionData` to store and manipulate single observations $(R_i,\Gamma_i)$ in the pre-training phase: +```julia +struct FrictionData + atoms + friction_tensor + friction_indices +end +``` +where +- `atoms` -- stores data of the atomic configuration ans is assumed to be of type `JuLIP.Atoms`, +- `friction_tensor` -- stores data on the friction tensor and is assumed to be of type`SparseMatrix{SMatrix{3, 3, T, 9}}`, where `T<:Float` and `Ti<:Int`. That is, the friction tensor is stored in the form of sparse matrices with $3 \times 3$-matrix valued block entries, +-`friction_indices` -- is a one-dimensional integer array, which contains all atom indices for which the friction tensor is defined. + +# Importing & Exporting Friction Data + +`ACEfriction.DataUtils` implements the function [`save_h5fdata(rdata::Vector{FrictionData}, filename::String)`](@ref) to save arrays of friction data to a custom costum formatted `hdf5` file, as well as the function [`load_h5fdata(filename::String)`](@ref) to load friction data from such costum formatted `hdf5` files. + + +# Custom HDF5 File Format for Friction Data + +The hierachical structure of such hdf5 files is as follows: + +Each observation $(R_i,\Gamma_i)$ is saved in a separate group in named by respective index $i$, i.e., we have the following groups on root level of the hdf5 file: +``` +├─ 📂 1 +├─ 📂 2 +├─ 📂 3 +│ : +├─ 📂 N_obs +``` +Within each of these groups, the data of the respective atomic configuration $R_i$ and friction tensor $\Gamma_i$ are saved in the subgroups `atoms` and `friction_tensor`, respectively: +``` +├─ 📂 i # Index of data point + ├─ 📂 atoms # Atom configuration data + │ ├─ 🔢 atypes + │ ├─ 🔢 cell + │ │ └─ 🏷️ column_major + │ ├─ 🔢 pbc + │ └─ 🔢 positions + │ └─ 🏷️ column_major + └─ 📂 friction_tensor # Friction tensor data + ├─ 🔢 ft_I + ├─ 🔢 ft_J + ├─ 🔢 ft_mask + └─ 🔢 ft_val, + └─ 🏷️ column_major +``` + +[Datasets](https://support.hdfgroup.org/documentation/hdf5/latest/_h5_d__u_g.html) in the group `atoms` store the equivalent information provided by the attributes `positions`, `numbers`, `cell`, and `pbc` of [atoms objects](https://wiki.fysik.dtu.dk/ase/ase/atoms.html), i.e., +an atomic configuration of N atoms is described by the following datasets contained in the group `atoms`: +- `atypes` -- A one-dimensional Integer dataset of length N. The ith entry corresponds to the atomic element number of the ith atom in the configuration. (Note: `types` corresponds to the atoms attribute `numbers` in the ase.) +- `cell` -- A two-dimensional Float64 dataset of size 3 x 3. +- `pbc` -- A one-dimensional Integer array of length 3 indicating the periodicity properties of the xyz dimension, e.g., `pbc = [1,0,0]` describes periodic boundary conditions in x dimension and non-periodic boundary conditions in the y and z dimensions. +- `positions` -- A two-dimensional Float64 dataset of size n N x 3. The ith column corresponds to the position of the ith atom in the configuration + +[Datasets](https://support.hdfgroup.org/documentation/hdf5/latest/_h5_d__u_g.html) in the group `friction_tensor` store the friction tensor as a N x N sparse matrix with 3x3 valued block entries, i.e., a friction tensor with m non-zero 3x3 blocks is stored +- `ft_I` -- A one-dimensional Integer dataset of length m specifying the column indices of non-zero block entries of the friction tensor. +- `ft_J` -- A one-dimensional Integer dataset of length m specifying the row indices of non-zero block entries of the friction tensor. +- `ft_val` -- A three-dimensional Float64 dataset of size m x 3 x 3 specifying the values of the non-zero 3 x 3 block entries of the friction tensor. For example, the 3 x 3 array `ft_val[k,:,:]` corresponds to the 3 x 3 block entry of the friction tensor with column index `ft_I[k]`, and row index `ft_J[k]`. +- `ft_mask` -- A one-dimensional Integer dataset/list containg the indices of atoms for which friction information is provided. + + +!!! note + All two or three-dimensional datasets in the groups `atoms` have an additional attribute `column_major`. If the hdf5 file is created in a language that stores matrices in column-major form (e.g., julia), this attribute must be set to 1 (True). If the hdf5 file is created in a language that stores matrices in column-major form (e.g., python), this attribute must be set to 0 (False). + + diff --git a/docs/src/fitting-eft.md b/docs/src/fitting-eft.md index d458df6..2fe1f1c 100644 --- a/docs/src/fitting-eft.md +++ b/docs/src/fitting-eft.md @@ -3,7 +3,7 @@ In this workflow example we demonstrate how `ACEfriction.jl` can be used to fit a simple 6 x 6 Electronic friction tensor modeling the non-adiabitic interactions of a hydrogen-atom on a copper surface. ## Load Electronic Friction Tensor Data -We first use the function [load_h5fdata]() to load the data of friction tensors from a [custom-formated]() hdf5 file and convert the data to the internal data format [FrictionData]. +We first use the function [load_h5fdata](@ref) to load the data of friction tensors from a [custom-formated](@ref costum-hdf5-format) hdf5 file and convert the data to the internal data format [FrictionData](@ref friction-data-representation). ```julia using ACEfriction # Load data @@ -12,8 +12,8 @@ rdata = ACEfriction.DataUtils.load_h5fdata( "./test/test-data-100.h5"); n_train = Int(ceil(.8 * length(rdata))) n_test = length(rdata) - n_train # Partition data into train and test set and convert the data -fdata = Dict("train" => FrictionData.(rdata[1:n_train]), - "test"=> FrictionData.(rdata[n_train+1:end])); +fdata = Dict("train" => rdata[1:n_train], + "test"=> rdata[n_train+1:end]); ``` ## Specify the Friction Model @@ -46,7 +46,7 @@ To train our model we first extract the parameters from the friction model, whic c=params(fm) ffm = FluxFrictionModel(c) ``` -Next, the function `flux_assemble` is used to prepare data for training. This includes evaluating the ACE-basis functions of the matrix models in `fm` on all configurations in the data set. Since the loss function of our model is quartic polynomial in the parameters, we don't need to reevaluate the ACE-basis functions at later stages of the training process. +Next, the function [`flux_assemble`]() is used to prepare data for training. This includes evaluating the ACE-basis functions of the matrix models in `fm` on all configurations in the data set. Since the loss function of our model is quartic polynomial in the parameters, we don't need to reevaluate the ACE-basis functions at later stages of the training process. ```julia flux_data = Dict( "train"=> flux_assemble(fdata["train"], fm, ffm; ), "test"=> flux_assemble(fdata["test"], fm, ffm)); diff --git a/docs/src/fitting-mbdpd.md b/docs/src/fitting-mbdpd.md index a0ef5b4..e0fdcb8 100644 --- a/docs/src/fitting-mbdpd.md +++ b/docs/src/fitting-mbdpd.md @@ -34,8 +34,8 @@ The following code loads training and test data comprised of particle configurat rdata_train = ACEfriction.DataUtils.load_h5fdata("./examples/data/dpd-train-x.h5"); rdata_test = ACEfriction.DataUtils.load_h5fdata("./examples/data/dpd-train-x.h5"); -fdata = Dict("train" => FrictionData.(rdata_train), - "test"=> FrictionData.(rdata_test)); +fdata = Dict("train" => rdata_train, + "test"=> rdata_test); (n_train, n_test) = length(fdata["train"]), length(fdata["test"]) ``` Here the training data is contains friction tensors of 50 configurations each comprised of 64 particles, and the test data contains friction tensors of 10 configurations each comprised of 216 particles. The underlying friction tensors were synthetically generated using the following simple friction model, which is a smooth version of the standard heuristic DPD friction models commonly used in simulations: diff --git a/docs/src/function-manual/ACEfriction.DataUtils.md b/docs/src/function-manual/ACEfriction.DataUtils.md new file mode 100644 index 0000000..68d5ca5 --- /dev/null +++ b/docs/src/function-manual/ACEfriction.DataUtils.md @@ -0,0 +1,85 @@ +```@meta +CurrentModule = ACEfriction.DataUtils +``` + +The submodule `ACEfriction.DataUtils.jl` provides structures to internally store friction data in `ACEfriction.jl` and functions to import and export friction data from and to custom formatted hdf5 files. + + +### [Friction Data representation](@id friction-data-representation) + +For pre-training storage and manipulation of friction data, +`ACEfriction.jl` implements the structure `FrictionData` to represent a single observation $(R_i,\Gamma_i)$ of an atomic configuration $R_i$ and corresponding friction tensor $\Gamma_i$: +```julia +struct FrictionData + atoms + friction_tensor + friction_indices +end +``` +where +- `atoms` -- stores data of the atomic configuration ans is assumed to be of type `JuLIP.Atoms`, +- `friction_tensor` -- stores data on the friction tensor and is assumed to be of type`SparseMatrix{SMatrix{3, 3, T, 9}}`, where `T<:Float` and `Ti<:Int`. That is, the friction tensor is stored in the form of sparse matrices with $3 \times 3$-matrix valued block entries, +- `friction_indices` -- is a one-dimensional integer array, which contains all atom indices for which the friction tensor is defined. + +### Importing & Exporting Friction Data + +`ACEfriction.DataUtils` implements the function [`save_h5fdata`](@ref) to save arrays of friction data to a [custom formatted](@ref costum-hdf5-format) `hdf5` file, as well as the function [`load_h5fdata`](@ref) to load friction data from such costum formatted `hdf5` files: + +```@docs +save_h5fdata +``` + +```@docs +load_h5fdata +``` + + +### [Custom HDF5 File Format for Friction Data](@id costum-hdf5-format) + +The hierachical structure of such hdf5 files is as follows: + +Each observation $(R_i,\Gamma_i)$ is saved in a separate group in named by respective index $i$, i.e., we have the following groups on root level of the hdf5 file: +``` +├─ 📂 1 +├─ 📂 2 +├─ 📂 3 +│ : +├─ 📂 N_obs +``` +Within each of these groups, the data of the respective atomic configuration $R_i$ and friction tensor $\Gamma_i$ are saved in the subgroups `atoms` and `friction_tensor`, respectively: +``` +├─ 📂 i # Index of data point + ├─ 📂 atoms # Atom configuration data + │ ├─ 🔢 atypes + │ ├─ 🔢 cell + │ │ └─ 🏷️ column_major + │ ├─ 🔢 pbc + │ └─ 🔢 positions + │ └─ 🏷️ column_major + └─ 📂 friction_tensor # Friction tensor data + ├─ 🔢 ft_I + ├─ 🔢 ft_J + ├─ 🔢 ft_mask + └─ 🔢 ft_val, + └─ 🏷️ column_major +``` + +[Datasets](https://support.hdfgroup.org/documentation/hdf5/latest/_h5_d__u_g.html) in the group `atoms` store the equivalent information provided by the attributes `positions`, `numbers`, `cell`, and `pbc` of [atoms objects](https://wiki.fysik.dtu.dk/ase/ase/atoms.html), i.e., +an atomic configuration of N atoms is described by the following datasets contained in the group `atoms`: +- `atypes` -- A one-dimensional Integer dataset of length N. The ith entry corresponds to the atomic element number of the ith atom in the configuration. (Note: `types` corresponds to the atoms attribute `numbers` in the ase.) +- `cell` -- A two-dimensional Float64 dataset of size 3 x 3. +- `pbc` -- A one-dimensional Integer array of length 3 indicating the periodicity properties of the xyz dimension, e.g., `pbc = [1,0,0]` describes periodic boundary conditions in x dimension and non-periodic boundary conditions in the y and z dimensions. +- `positions` -- A two-dimensional Float64 dataset of size n N x 3. The ith column corresponds to the position of the ith atom in the configuration + +[Datasets](https://support.hdfgroup.org/documentation/hdf5/latest/_h5_d__u_g.html) in the group `friction_tensor` store the friction tensor as a N x N sparse matrix with 3x3 valued block entries, i.e., a friction tensor with m non-zero 3x3 blocks is stored +- `ft_I` -- A one-dimensional Integer dataset of length m specifying the column indices of non-zero block entries of the friction tensor. +- `ft_J` -- A one-dimensional Integer dataset of length m specifying the row indices of non-zero block entries of the friction tensor. +- `ft_val` -- A three-dimensional Float64 dataset of size m x 3 x 3 specifying the values of the non-zero 3 x 3 block entries of the friction tensor. For example, the 3 x 3 array `ft_val[k,:,:]` corresponds to the 3 x 3 block entry of the friction tensor with column index `ft_I[k]`, and row index `ft_J[k]`. +- `ft_mask` -- A one-dimensional Integer dataset/list containg the indices of atoms for which friction information is provided. + + +!!! warning + All two or three-dimensional datasets in the groups `atoms` and `friction_tensor` have an additional attribute `column_major`. If the hdf5 file is created in a language that stores matrices in column-major form (e.g., julia), this attribute must be set to 1 (True). If the hdf5 file is created in a language that stores matrices in column-major form (e.g., python), this attribute must be set to 0 (False). + + + diff --git a/docs/src/function-manual/ACEfriction.FrictionFit.md b/docs/src/function-manual/ACEfriction.FrictionFit.md new file mode 100644 index 0000000..20feac1 --- /dev/null +++ b/docs/src/function-manual/ACEfriction.FrictionFit.md @@ -0,0 +1,23 @@ +```@meta +CurrentModule = ACEfriction.FrictionFit +``` + +```@docs +flux_assemble +``` +```@docs +FluxFrictionModel +``` + +```@docs +get_ids +``` +```@docs +l2_loss +``` +```@docs +weighted_l2_loss +``` +```@docs +weighted_l1_loss +``` \ No newline at end of file diff --git a/docs/src/function-manual.md b/docs/src/function-manual/ACEfriction.FrictionModels.md similarity index 87% rename from docs/src/function-manual.md rename to docs/src/function-manual/ACEfriction.FrictionModels.md index e7625b3..fa0b831 100644 --- a/docs/src/function-manual.md +++ b/docs/src/function-manual/ACEfriction.FrictionModels.md @@ -1,7 +1,3 @@ -# Function Manual (@id Function-Manual) - -## ACEfriction.FrictionModels.jl - ```@meta CurrentModule = ACEfriction.FrictionModels ``` @@ -47,8 +43,3 @@ set_params! ```@docs set_zero! ``` - -## ACEfriction.MatrixModels.jl - - -## ACEfriction.FrictionFit.jl \ No newline at end of file diff --git a/docs/src/function-manual/ACEfriction.MatrixModels.md b/docs/src/function-manual/ACEfriction.MatrixModels.md new file mode 100644 index 0000000..bad57c1 --- /dev/null +++ b/docs/src/function-manual/ACEfriction.MatrixModels.md @@ -0,0 +1,16 @@ +```@meta +CurrentModule = ACEfriction.MatrixModels +``` + +```@docs +RWCMatrixModel +``` +```@docs +PWCMatrixModel +``` +```@docs +OnsiteOnlyMatrixModel +``` +```@docs +mbdpd_matrixmodel +``` \ No newline at end of file diff --git a/docs/src/index.md b/docs/src/index.md index f21bf46..9d3664d 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -2,7 +2,7 @@ The julia package `ACEfriction.jl` facilitates simulation and machine learning of configuration-dependent friction tensor models from data. The models are based on an equivariant Atomic Cluster Expansion (ACE) and, as such, are computationally highly efficient and size transferable. The underlying framework of model construction is described in detail in [Sachs et al., (2024)](@ref ACEfriction-paper). -For a quick start, we recommend reading the [Installation Instructions](installation.md) and the [Overview](overview.md) section, followed by the [Workflow Examples](fitting-eft.md). Detailed documentation of front-end-facing functions can be found in the [Function Manual](function-manual.md). +For a quick start, we recommend reading the [Installation Instructions](installation.md) and the [Overview](overview.md) section, followed by the [Workflow Examples](fitting-eft.md). Detailed documentation of front-end-facing functions can be found in the Function Manual. ### [References](@id ACEfriction-paper) diff --git a/examples/fit-ACEFrictionModel-new-ac.jl b/examples/fit-ACEFrictionModel-new-ac.jl index ee48882..e9b05e6 100644 --- a/examples/fit-ACEFrictionModel-new-ac.jl +++ b/examples/fit-ACEFrictionModel-new-ac.jl @@ -93,6 +93,7 @@ set_params!(fm, params(ffm)) at = fdata["test"][1].atoms @time Gamma(fm, at) +G = Gamma(fm, at) @time Σ = Sigma(fm, at) @time Gamma(fm, Σ) @time randf(fm, Σ) diff --git a/src/ACEfriction.jl b/src/ACEfriction.jl index b473dc4..d28da7c 100644 --- a/src/ACEfriction.jl +++ b/src/ACEfriction.jl @@ -9,7 +9,6 @@ include("./atomcutoffs.jl") include("./matrixmodels/matrixmodels.jl") include("./frictionmodels.jl") include("./frictionfit/frictionfit.jl") - include("./matrixmodelsutils.jl") import ACEfriction.FrictionModels: FrictionModel, Gamma, Sigma diff --git a/src/datautils.jl b/src/datautils.jl index aeeeb56..f6f90ea 100644 --- a/src/datautils.jl +++ b/src/datautils.jl @@ -9,6 +9,18 @@ export FrictionData, BlockDenseArray export save_h5fdata, load_h5fdata using HDF5 +struct FrictionData + atoms + friction_tensor + friction_indices +end + +function FrictionData(d::NamedTuple{(:at, :friction_tensor, :friction_indices)}) + return FrictionData(d.at, d.friction_tensor, d.friction_indices) +end +# function FrictionData(d::NamedTuple{(:at, :friction_tensor, :friction_indices,:friction_indices_ref)}) +# return FrictionData(d.at, d.friction_tensor, d.friction_indices, d.friction_tensor_ref) +# end function _array2svector(x::Array{T,2}) where {T} return [ SVector{3}(x[i,:]) for i in 1:size(x)[1] ] @@ -26,30 +38,21 @@ function _svector2array(c_vec::Vector{SVector{N_rep, T}}) where {N_rep,T<:Number return c_matrix end -""" - save_h5fdata(rdata, filename ) - -Saves a friction tensor dataset with the HDF5 file structure expected by `load_h5fdata`, for use in training scripts. - -# Arguments -## `rdata`: -A `Vector{NamedTuple}` containing the training data in the following format. - -``` -( - at = ::JuLIP.Atoms, - friction_tensor = sparse(::Matrix), - friction_indices = ::Vector{Int} -) -``` - -## `filename` +""" + save_h5fdata(rdata::Vector{FrictionData}, filename::String ) -Name of the file to save to (including h5 extension). +Saves a friction tensor data in a costum formatted hdf5 file. +### Arguments +- `rdata` : Vector{FrictionData} : + A vector of friction data entries. Each entry is a structure of type `Frictiondata` with the following fields: + - `at` : JuLIP.Atoms : Atoms object containing the atomic positions, cell, and periodic boundary conditions. + - `friction_tensor` : SparseMatrix{SMatrix{3,3,Float64,9}} : Sparse matrix representation of the friction tensor. + - `friction_indices` : Vector{Int} : Indices of the atoms for which the friction tensor is defined. +- `filename` : String : Name of the file to save to (including h5 extension). """ -function save_h5fdata(rdata, filename ) +function save_h5fdata(rdata::Vector{FrictionData}, filename::String ) fid = h5open(filename, "w") try # iterate over each data entry @@ -58,13 +61,13 @@ function save_h5fdata(rdata, filename ) g = create_group(fid, "$i") # write atoms data ag = create_group(g, "atoms") - dset_pos = create_dataset(ag, "positions", Float64, (length(d.at.X), 3)) - for (k,x) in enumerate(d.at.X) + dset_pos = create_dataset(ag, "positions", Float64, (length(d.atoms.X), 3)) + for (k,x) in enumerate(d.atoms.X) dset_pos[k,:] = x end - write(ag, "atypes", Int.(d.at.Z)) - write(ag, "cell", Matrix(d.at.cell)) - write(ag, "pbc", Array(d.at.pbc)) + write(ag, "atypes", Int.(d.atoms.Z)) + write(ag, "cell", Matrix(d.atoms.cell)) + write(ag, "pbc", Array(d.atoms.pbc)) # write friction data fg = create_group(g, "friction_tensor") (I,J,V) = findnz(d.friction_tensor) @@ -115,7 +118,7 @@ function _hdf52Atoms( ag::HDF5.Group ) end function _hdf52ft( ftg::HDF5.Group ) - local ft_val + local ft_val try if Bool(read_attribute(ftg["ft_val"],"column_major")) == true ft_val = read(ftg["ft_val"]) @@ -131,7 +134,23 @@ function _hdf52ft( ftg::HDF5.Group ) return (friction_tensor = spft, mask = ft_mask) end -function load_h5fdata(filename) + +""" + load_h5fdata(filename::String) + +Loads a friction tensor data from a costum formatted hdf5 file. + +### Arguments +- `filename` : String : Name of the file to load from (including h5 extension). + +### Returns +- `rdata` : Vector{FrictionData} : + A vector of friction data entries. Each entry is a structure of type `Frictiondata` with the following fields: + - `at` : JuLIP.Atoms : Atoms object containing the atomic positions, cell, and periodic boundary conditions. + - `friction_tensor` : SparseMatrix{SMatrix{3,3,Float64,9}} : Sparse matrix representation of the friction tensor. + - `friction_indices` : Vector{Int} : Indices of the atoms for which the friction tensor is defined. +""" +function load_h5fdata(filename::String) fid = h5open(filename, "r") N_data = read_attribute(fid, "N_data") rdata = @showprogress [begin @@ -141,7 +160,7 @@ function load_h5fdata(filename) end for i=1:N_data] HDF5.close(fid) - return rdata + return FrictionData.(rdata) end """ @@ -162,23 +181,7 @@ function BlockDenseArray(full_tensor::Matrix; indices=1:size(full_tensor,1)) return BlockDenseMatrix(full_tensor[indices,indices], indices) end -struct FrictionData{A} - atoms::Atoms - friction_tensor::A - friction_indices - friction_tensor_ref -end -function FrictionData(d::NamedTuple{(:at, :friction_tensor, :friction_indices)}) - return FrictionData(d.at, d.friction_tensor, d.friction_indices, nothing) -end -function FrictionData(d::NamedTuple{(:at, :friction_tensor, :friction_indices,:friction_indices_ref)}) - return FrictionData(d.at, d.friction_tensor, d.friction_indices, d.friction_tensor_ref) -end -# function FrictionData(atoms::Atoms, friction_tensor, friction_indices; -# friction_tensor_ref=nothing) -# return FrictionData(atoms, friction_tensor, friction_indices, weights, friction_tensor_ref) -# end end diff --git a/src/frictionfit/fdatautils.jl b/src/frictionfit/fdatautils.jl index 216a0f5..a86eca0 100644 --- a/src/frictionfit/fdatautils.jl +++ b/src/frictionfit/fdatautils.jl @@ -1,5 +1,6 @@ + """ - flux_assemble(data::Array{DATA}, fm::FrictionModel, ffm::FluxFrictionModel) where {DATA<:FrictionData} + flux_assemble(fdata::Array{DATA}, fm::FrictionModel, ffm::FluxFrictionModel; weights = Dict("observations" => ones(length(fdata)), "diag" => 2.0, "sub_diag" => 1.0, "off_diag"=>1.0)) where {DATA<:FrictionData} Converts FrictionData into a format that can be used for training with ffm::FluxFrictionModel @@ -105,11 +106,7 @@ end """ function _flux_data(d::FrictionData,fm::FrictionModel, transforms::NamedTuple, W, join_sites=true) # TODO: in-place data manipulations - if d.friction_tensor_ref === nothing - friction_tensor = _tensor_Gamma(d.friction_tensor,d.friction_indices) - else - friction_tensor = _tensor_Gamma(d.friction_tensor-d.friction_tensor_ref,d.friction_indices) - end + friction_tensor = _tensor_Gamma(d.friction_tensor,d.friction_indices) Tfm = Tuple(typeof(mo) for mo in values(fm.matrixmodels)) BB = basis(fm, d.atoms; join_sites=join_sites) BB = Tuple(_tensor_basis(transform_basis(B,trans),d.friction_indices, tfm) for (B,tfm,trans) in zip(BB,Tfm,transforms)) diff --git a/src/frictionfit/fluxmodels.jl b/src/frictionfit/fluxmodels.jl index 88d05c7..ca25852 100644 --- a/src/frictionfit/fluxmodels.jl +++ b/src/frictionfit/fluxmodels.jl @@ -92,6 +92,16 @@ function _add_default_transforms(transforms::NamedTuple, model_ids) ) for s in model_ids ) end + + +""" + FluxFrictionModel(c::NamedTuple{model_ids}) where {model_ids} + +Create a FluxFrictionModel with the parameters `c` and the model ids `model_ids`. + +### Arguments +- `c::NamedTuple{model_ids}` : The parameters of the model. +""" function FluxFrictionModel(c::NamedTuple{model_ids}; transforms::NamedTuple=NamedTuple()) where {model_ids} transform_filtered = _add_default_transforms(transforms, model_ids) #return FluxFrictionModel( map(cc->reinterpret(SVector{Vector{Float64}}, cc), transform_params(Tuple(c),transform_filtered)), model_ids, transform_filtered) @@ -129,6 +139,16 @@ end # end # end +""" + set_params!(m::FluxFrictionModel; sigma=1E-8, model_ids::Array{Symbol}=Symbol[]) + +Randomizes the parameters of the model `m`. + +### Arguments +- `m::FluxFrictionModel`: The model to set the parameters of. +- `sigma::Float64=1E-8`: The standard deviation of the random values. +- `model_ids::Array{Symbol}=[]`: The ids of the models to set the parameters of. If empty, all the parameters are set. +""" function set_params!(m::FluxFrictionModel; sigma=1E-8, model_ids::Array{Symbol}=Symbol[]) model_ids = (isempty(model_ids) ? get_ids(m) : model_ids) for (v,s) in zip(m.c,m.model_ids) @@ -139,6 +159,15 @@ function set_params!(m::FluxFrictionModel; sigma=1E-8, model_ids::Array{Symbol}= end end +""" + set_params!(m::FluxFrictionModel, c_new::NamedTuple) + +Set the parameters of the model to the values in `c_new`. + +### Arguments +- `m::FluxFrictionModel`: The model to set the parameters of. +- `c_new::NamedTuple`: The new parameters. The keys of `c_new` should be a subset of `m.model_ids`. +""" function set_params!(m::FluxFrictionModel, c_new::NamedTuple) for (v,s) in zip(m.c, m.model_ids) if s in keys(c_new) @@ -146,13 +175,23 @@ function set_params!(m::FluxFrictionModel, c_new::NamedTuple) end end end +""" + get_ids(m::FluxFrictionModel) +Return the model ids of the model `m`. +""" get_ids(m::FluxFrictionModel) = m.model_ids (m::FluxFrictionModel)(B, Tfm) = _Gamma(B, m.c, Tfm) Flux.@functor FluxFrictionModel (c,) Flux.trainable(m::FluxFrictionModel) = (c=m.c,) + +""" + params(m::FluxFrictionModel; transformed=true) + +Return the parameters of the model `m`. +""" params(m::FluxFrictionModel; transformed=true) = NamedTuple{m.model_ids}(transformed ? rev_transform_params(m.c,m.transforms) : m.c ) get_transform(m::FluxFrictionModel) = NamedTuple{m.model_ids}(m.transforms) @@ -160,17 +199,37 @@ function _l2(Γ_fit::Array{T,4},Γ_true::Array{T,4}) where {T<:Number} @tullio err:= (Γ_fit[d1,d2,i,j]- Γ_true[d1,d2,i,j])^2 return err end + +""" + l2_loss(fm, data) + +Compute the L2 loss of the model `fm` on the data `data`. +""" l2_loss(fm, data) = sum(_l2(fm(d.B, d.Tfm), d.friction_tensor) for d in data) function _weighted_l2(Γ_fit::Array{T,4},Γ_true::Array{T,4},W::Array{T,4}) where {T<:Number} @tullio err:= W[d1,d2,i,j] * (Γ_fit[d1,d2,i,j]- Γ_true[d1,d2,i,j])^2 return err end -weighted_l2_loss(fm, data) = sum(_weighted_l2(fm(d.B, d.Tfm), d.friction_tensor, d.W) for d in data) + + +""" + weighted_l2_loss(fm, data) + +Compute the weighted L2 loss of the model `fm` on the data `data`. +""" +function weighted_l2_loss(fm, data) + sum(_weighted_l2(fm(d.B, d.Tfm), d.friction_tensor, d.W) for d in data) +end function _weighted_l1(Γ_fit::Array{T,4},Γ_true::Array{T,4},W::Array{T,4}) where {T<:Number} @tullio err:= W[d1,d2,i,j] * abs(Γ_fit[d1,d2,i,j]- Γ_true[d1,d2,i,j]) return err end +""" + weighted_l1_loss(fm, data) + +Compute the weighted L1 loss of the model `fm` on the data `data`. +""" weighted_l1_loss(fm, data) = sum(_weighted_l1(fm(d.B, d.Tfm), d.friction_tensor, d.W) for d in data) \ No newline at end of file diff --git a/src/matrixmodels/matrixmodels.jl b/src/matrixmodels/matrixmodels.jl index 83bde2f..6a91fe9 100644 --- a/src/matrixmodels/matrixmodels.jl +++ b/src/matrixmodels/matrixmodels.jl @@ -703,4 +703,5 @@ include("./pwcmatrixmodels.jl") # Omsite-only matrix models: include("./onsiteonlymatrixmodels.jl") + end \ No newline at end of file diff --git a/src/matrixmodels/pwcmatrixmodels.jl b/src/matrixmodels/pwcmatrixmodels.jl index 5e0f995..aa89806 100644 --- a/src/matrixmodels/pwcmatrixmodels.jl +++ b/src/matrixmodels/pwcmatrixmodels.jl @@ -1,6 +1,3 @@ -""" - -""" struct PWCMatrixModel{O3S,CUTOFF,Z2S,SC} <: MatrixModel{O3S} offsite::OffSiteModels{O3S,Z2S,CUTOFF} where {Z2S, CUTOFF} n_rep::Int diff --git a/src/matrixmodelsutils.jl b/src/matrixmodelsutils.jl index f576bc0..69cc76b 100644 --- a/src/matrixmodelsutils.jl +++ b/src/matrixmodelsutils.jl @@ -11,6 +11,7 @@ export RWCMatrixModel, PWCMatrixModel, OnsiteOnlyMatrixModel, mbdpd_matrixmodel # Outer convenience constructors for subtypes of MatrixModels + """ function RWCMatrixModel(property, species_friction, species_env; maxorder=2, @@ -36,7 +37,6 @@ Creates a matrix model with row-wise coupling. By default, this model evaluates - `n_rep` -- the number of matrix blocks evaluated per atom pair. - `species_substrat` -- a list of chemical element types. At least one atom of such element types must be within the pair-environemt of two friction-feeling atoms i,j in order for the matrix-block ``\\Sigma_{ij}`` to be non-zero. """ -# #- `rcut` -- For row-wise coupled matrix models, the pair environment of the atom pair i,j is by default defined as the set of atoms within a spherical cutoff of radius `rcut` around the atom i. function RWCMatrixModel(property, species_friction, species_env; maxorder=2, maxdeg=5, @@ -54,6 +54,7 @@ function RWCMatrixModel(property, species_friction, species_env; bond_weight = 1.0, id=nothing ) + # #- `rcut` -- For row-wise coupled matrix models, the pair environment of the atom pair i,j is by default defined as the set of atoms within a spherical cutoff of radius `rcut` around the atom i. return RWCMatrixModel(property, species_friction, species_env, evalcenter; n_rep = n_rep, species_substrat = species_substrat, @@ -231,7 +232,7 @@ end ) -Create a matrix model for a momentum-preserving friction tensors suitable for the simulation of Dissipative Particle Dynamics. The model is a particular parametrization of a pair-wise coupled matrix model. +Create a matrix model for a momentum-preserving friction model for the simulation of Dissipative Particle Dynamics. The model is a particular parametrization of a pair-wise coupled matrix model. This model evaluates blocks ``\\Sigma_{ij}`` as a function of ellipoid-shaped pair environments centered at the midpoints of the positions of atoms i.j. diff --git a/test/test_IO_data.jl b/test/test_IO_data.jl index b336619..e201b5f 100644 --- a/test/test_IO_data.jl +++ b/test/test_IO_data.jl @@ -14,6 +14,6 @@ save_h5fdata(data1,filename2); data2= load_h5fdata(filename2); rm(filename2) -@test all([all([getfield(d1.at,f) == getfield(d2.at,f) && -typeof(getfield(d1.at,f)) == typeof(getfield(d2.at,f)) for f in fieldnames(typeof(d1.at))]) +@test all([all([getfield(d1.atoms,f) == getfield(d2.atoms,f) && +typeof(getfield(d1.atoms,f)) == typeof(getfield(d2.atoms,f)) for f in fieldnames(typeof(d1.atoms))]) for (d1,d2) in zip(data1,data2)]) \ No newline at end of file diff --git a/test/test_ac_model_fit.jl b/test/test_ac_model_fit.jl index 1ad77a9..279a2c3 100644 --- a/test/test_ac_model_fit.jl +++ b/test/test_ac_model_fit.jl @@ -74,8 +74,8 @@ shuffle!(rng, rdata) n_train = Int(ceil(.8 * length(rdata))) n_test = length(rdata) - n_train -fdata = Dict("train" => FrictionData.(rdata[1:n_train]), - "test"=> FrictionData.(rdata[n_train+1:end])); +fdata = Dict("train" => rdata[1:n_train], + "test"=> rdata[n_train+1:end]); @info "Fit RWC friction model" c = params(fm) diff --git a/test/test_pwcec_model_fit.jl b/test/test_pwcec_model_fit.jl index a278e32..ce3e781 100644 --- a/test/test_pwcec_model_fit.jl +++ b/test/test_pwcec_model_fit.jl @@ -83,8 +83,8 @@ shuffle!(rng, rdata) n_train = Int(ceil(.8 * length(rdata))) n_test = length(rdata) - n_train -fdata = Dict("train" => FrictionData.(rdata[1:n_train]), - "test"=> FrictionData.(rdata[n_train+1:end])); +fdata = Dict("train" => rdata[1:n_train], + "test"=> rdata[n_train+1:end]); @info "Fit friction model" diff --git a/test/test_pwcsc_model_fit.jl b/test/test_pwcsc_model_fit.jl index 0f9841d..4b9cb59 100644 --- a/test/test_pwcsc_model_fit.jl +++ b/test/test_pwcsc_model_fit.jl @@ -82,8 +82,8 @@ shuffle!(rng, rdata) n_train = Int(ceil(.8 * length(rdata))) n_test = length(rdata) - n_train -fdata = Dict("train" => FrictionData.(rdata[1:n_train]), - "test"=> FrictionData.(rdata[n_train+1:end])); +fdata = Dict("train" => rdata[1:n_train], + "test"=> rdata[n_train+1:end]); @info "Fit friction model"