Add Subsystems to apply models to selected atoms only (#39)

* Extend Harmonic for multiple particles and oscillators for tests

* Add Subsystem and CompositeModel types

A CompositeModel combines multiple models assigned to determine parts
of the total system, e.g. different friction processes for different
parts of the system.

The Subsystem type is the building block of a CompositeModel, combining
an arbitrary NQCModels.Model with indices determining which particles
in the system it applies to.

* Made ElectronicFrictionProviders a subtype of Model

This is so they are included as possible inputs for a Subsystem.

* Load CompositeModels

* CompositeModels are abstract Model sustypes

* Forgot to add dofs functions for CompositeModels

* Subsystem and CompositeModel have custom Base.show routines

* Subsystems should also accept colon indices

* Subsystem index logic now less type-constrained

* Diagnosing a StackOverflow with ASE models

* Possibly fixed type problems causing a stack overflow

Or not

Another try

Debug

Tried to break the stack overflow

* ElectronicFrictionProviders get dofs and ndofs too

* Attempt to fix friction matrix generation for subsystems

Need to call friction! and derivative! for subsystems and full R

I am very confused

* Made ASE interface aware of FixAtoms constraints

* Remove accidentally added ase PyCall extension with duplicate functionality

* Remove ASEconvert as no longer needed

* Ensure ase is available through CondaPkg

* Update docs

* Remove debug info

* Add unit test and fix small bug

* Naming inconsistency between dofs and ndofs

* CompositeModel derivative fixed.

* Version bump #39

.gitignore

@@ -26,3 +26,6 @@ Manifest.toml
 .DS_STORE
 .CondaPkg/
 .CondaPkg/*
+
+# IDE files
+.vscode/

CondaPkg.toml

@@ -0,0 +1,4 @@
+[deps]
+numpy = "<2" # Fix below 2 for now
+ase = ""
+

Project.toml

@@ -1,10 +1,9 @@
 name = "NQCModels"
 uuid = "c814dc9f-a51f-4eaf-877f-82eda4edad48"
 authors = ["James Gardner <[email protected]>"]
-version = "0.9.0"
+version = "0.9.1"
 
 [deps]
-ASEconvert = "3da9722f-58c2-4165-81be-b4d7253e8fd2"
 AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -32,7 +31,7 @@ Requires = "1"
 StaticArrays = "1"
 Unitful = "1"
 UnitfulAtomic = "1"
-julia = "1.7"
+julia = "1.9"
 
 [extras]
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
@@ -44,12 +43,4 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = [
-    "Plots",
-    "PyCall",
-    "SafeTestsets",
-    "Test",
-    "FiniteDiff",
-    "JuLIP",
-    "PythonCall",
-]
+test = ["Plots", "PyCall", "SafeTestsets", "Test", "FiniteDiff", "JuLIP", "PythonCall"]

src/NQCModels.jl

@@ -148,4 +148,6 @@ include("diabatic/DiabaticModels.jl")
 
 include("plot.jl")
 
+include("subsystems.jl")
+
 end # module

src/adiabatic/harmonic.jl

@@ -4,6 +4,10 @@
 
 Adiabatic harmonic potential. ``V(x) = mω^2(x-r₀)^2 / 2``
 
+m, ω, r₀ are the mass, frequency, and equilibrium position respectively and can be supplied as 
+numbers or Matrices to create a compound model for multiple particles. 
+
+
 ```jldoctest
 julia> using Symbolics;
 
@@ -29,9 +33,9 @@ end
 NQCModels.ndofs(harmonic::Harmonic) = harmonic.dofs
 
 function NQCModels.potential(model::Harmonic, R::AbstractMatrix)
-    return sum(0.5 * model.m* model.ω^2 .* (R .- model.r₀) .^2)
+    return sum(@. 0.5 * model.m* model.ω^2 * (R - model.r₀) ^2)
 end
 
 function NQCModels.derivative!(model::Harmonic, D::AbstractMatrix, R::AbstractMatrix) 
-    D .= model.m* model.ω^2 .* (R .- model.r₀)
+    @. D = model.m * model.ω^2 * (R - model.r₀)
 end

src/friction/FrictionModels.jl

@@ -40,6 +40,7 @@ function friction! end
 Obtain the friction for the current position `R`.
 
 This is an allocating version of `friction!`.
+
 """
 function friction(model::AdiabaticFrictionModel, R)
     F = zero_friction(model, R)
@@ -49,6 +50,9 @@ end
 
 zero_friction(::AdiabaticFrictionModel, R) = zeros(eltype(R), length(R), length(R))
 
+NQCModels.dofs(model::ElectronicFrictionProvider) = 1:model.ndofs
+NQCModels.ndofs(model::ElectronicFrictionProvider) = model.ndofs
+
 include("composite_friction_model.jl")
 export CompositeFrictionModel
 include("ase_friction_interface.jl")

src/subsystems.jl

@@ -0,0 +1,152 @@
+
+export Subsystem, CompositeModel
+using .FrictionModels
+using .AdiabaticModels
+
+"""
+Subsystem(M, indices)
+
+A subsystem is a Model which only applies to a subset of the degrees of freedom of the original model. 
+
+**When combined in a CompositeModel**, `potential()`, `derivative!()` and `friction!()` will be sourced from the respective Subsystems. 
+
+Calling `potential()`, `derivative!()`, or `friction!()` on a subsystem directly will output the respective values **for the entire system**. 
+
+The Model specified will be supplied with the positions of the entire system for evaluation. 
+
+"""
+struct Subsystem{M<:Union{Model, FrictionModels.ElectronicFrictionProvider}}
+	model::M
+	indices
+end
+
+function Base.show(io::IO, subsystem::Subsystem)
+    print(io, "Subsystem:\n\t🏎️ $(subsystem.model)\n\t🔢 $(subsystem.indices)\n")
+end
+
+function Subsystem(model, indices=:)
+	# Convert indices to a Vector{Int} or : for consistency
+	if isa(indices, Int)
+		indices = [indices:indices]
+	elseif isa(indices, UnitRange{Int})
+		indices = collect(indices)
+	end
+	Subsystem(model, indices)
+end
+
+# Passthrough functions to Model functions
+potential(subsystem::Subsystem, R::AbstractMatrix) = potential(subsystem.model, R)
+derivative(subsystem::Subsystem, R::AbstractMatrix) = derivative(subsystem.model, R)
+derivative!(subsystem::Subsystem, D::AbstractMatrix, R::AbstractMatrix) = derivative!(subsystem.model, D, R)
+FrictionModels.friction(subsystem::Subsystem, R::AbstractMatrix) = friction(subsystem.model, R)
+FrictionModels.friction!(subsystem::Subsystem, F::AbstractMatrix, R::AbstractMatrix) = friction!(subsystem.model, F, R)
+dofs(subsystem::Subsystem) = dofs(subsystem.model)
+ndofs(subsystem::Subsystem) = ndofs(subsystem.model)
+
+"""
+CompositeModel(Subsystems...)
+
+A CompositeModel is composed of multiple Subsystems, creating an effective model which evaluates each Subsystem for its respective indices. 
+"""
+struct CompositeModel{S<:Vector{<:Subsystem}, D<:Int} <: AdiabaticModels.AdiabaticModel
+	subsystems::S
+	ndofs::D
+end
+
+function Base.show(io::IO, model::CompositeModel)
+    print(io, "CompositeModel with subsystems:\n", [system for system in model.subsystems]...)
+end
+
+"""
+    CompositeModel(subsystems::Subsystem...)
+
+Combine multiple Subsystems into a single model to be handled by NQCDynamics.jl in simulations. 
+Any calls made to `potential`, `derivative` and `friction` will apply each subsystem's model to the respective atoms while ignoring any other atoms. 
+
+Some checks are made to ensure each atom is affected by a model and that each model is applied over the same degrees of freedom, but no other sanity checks are made. 
+"""
+CompositeModel(subsystems::Subsystem...) = CompositeModel(check_models(subsystems...)...) # Check subsystems are a valid combination
+
+get_friction_models(system::Vector{<:Subsystem}) = @view system[findall(x->isa(x.model, FrictionModels.ElectronicFrictionProvider), system)]
+get_friction_models(system::CompositeModel) = get_friction_models(system.subsystems)
+get_pes_models(system::Vector{<:Subsystem}) = @view system[findall(x->isa(x.model, AdiabaticModels.AdiabaticModel) || isa(x.model, DiabaticModels.DiabaticModel), system)]
+get_pes_models(system::CompositeModel) = get_pes_models(system.subsystems)
+
+dofs(system::CompositeModel) = 1:system.ndofs
+ndofs(system::CompositeModel) = system.ndofs
+
+"""
+Subsystem combination logic - We only want to allow combination of subsystems:
+
+? 1. with the same number of degrees of freedom
+
+2. without overlapping indices (Build a different type of CompositeModel) to handle these cases separately. 
+"""
+function check_models(subsystems::Subsystem...)
+	systems=vcat(subsystems...)
+	# Check unique assignment of potential and derivative to each atom index
+	pes_models = get_pes_models(systems)
+	pes_model_indices = vcat([subsystem.indices for subsystem in pes_models]...)
+	if length(unique(pes_model_indices)) != length(pes_model_indices)
+		error("Overlapping indices detected for the assignment of potential energy surfaces.")
+	end
+	# Check for unique assignment of friction to each atom index
+	model_has_friction=systems[findall(x->isa(x.model, FrictionModels.ElectronicFrictionProvider), systems)]
+	friction_indices = vcat([subsystem.indices for subsystem in model_has_friction]...)
+	if length(unique(friction_indices)) != length(friction_indices)
+		error("Overlapping indices detected for the assignment of friction models.")
+	end
+
+	# Check for different numbers of degrees of freedom in each subsystem
+	dofs = [ndofs(subsystem.model) for subsystem in subsystems]
+	if length(unique(dofs)) != 1
+		error("Subsystems must have the same number of degrees of freedom.")
+	end
+	return systems, unique(dofs)[1]
+end
+
+# Subsystem evaluation of model functions
+function potential(system::CompositeModel, R::AbstractMatrix)
+	pes_models=get_pes_models(system)
+	total_potential_energy=potential(pes_models[1], R)
+	for subsystem in pes_models[2:end]
+		total_potential_energy+=potential(subsystem, R)
+	end
+	return total_potential_energy
+end
+
+function derivative!(system::CompositeModel, D::AbstractMatrix, R::AbstractMatrix)
+	for subsystem in get_pes_models(system)
+		@debug "Accessing D[$(dofs(subsystem)), $(subsystem.indices)]"
+		subsystem_derivative = derivative(subsystem, R)
+		D[dofs(subsystem), subsystem.indices] .= subsystem_derivative[dofs(subsystem), subsystem.indices]
+	end
+end
+
+function derivative(system::CompositeModel, R::AbstractMatrix)
+	total_derivative=zero(R)
+	derivative!(system, total_derivative, R)
+	return total_derivative
+end
+
+function FrictionModels.friction!(system::CompositeModel, F::AbstractMatrix, R::AbstractMatrix)
+	for subsystem in get_friction_models(system)
+		if subsystem.indices == Colon()
+			eft_indices=vcat([[(j-1)*ndofs(subsystem.model)+i for i in dofs(subsystem.model)] for j in 1:size(R,2)]...) # Size of friction tensor from positions if applying friction to entire system. 
+		else
+			eft_indices=vcat([[(j-1)*ndofs(subsystem.model)+i for i in dofs(subsystem.model)] for j in subsystem.indices]...)
+		end
+		FrictionModels.friction!(subsystem, view(F, eft_indices, eft_indices), R)
+	end
+end
+
+function FrictionModels.friction(system::CompositeModel, R::AbstractMatrix)
+	F=zeros(eltype(R), length(R), length(R))
+	FrictionModels.friction!(system, F, R)
+	return F
+end
+
+
+
+
+

test/runtests.jl

@@ -4,6 +4,27 @@ using NQCModels
 using LinearAlgebra
 using SafeTestsets
 
+const GROUP = get(ENV, "GROUP", "All")
+
+if GROUP == "All" || GROUP == "Composite"
+    @testset "CompositeModel" begin
+        osc_1 = Harmonic(; ω = 1.0)
+        osc_2 = Harmonic(; ω = 2.0)
+        composite = CompositeModel(
+            Subsystem(osc_1, 1:1),
+            Subsystem(osc_2, 2:2)
+        )
+        positions = [3.45 1.23]
+        combined_derivative = hcat(
+            NQCModels.derivative(osc_1, positions[:, 1:1]),
+            NQCModels.derivative(osc_2, positions[:, 2:2]),
+        )
+        combined_derivative_values = zero(positions)
+        NQCModels.derivative!(composite, combined_derivative_values, positions)
+        @test combined_derivative == combined_derivative_values
+    end
+end
+
 @time @safetestset "Wide band bath discretisations" begin
     include("wide_band_bath_discretisations.jl")
 end
@@ -16,12 +37,14 @@ end
 
 include("test_utils.jl")
 
-@time @safetestset "ASE with PythonCall.jl" begin
-    include("ase_pythoncall.jl")
-end
-
-@time @safetestset "ASE with PyCall.jl" begin
-    include("ase_pycall.jl")
+if GROUP =="All" || GROUP == "ASE"
+    @time @safetestset "ASE with PythonCall.jl" begin
+        include("ase_pythoncall.jl")
+    end
+
+    @time @safetestset "ASE with PyCall.jl" begin
+        include("ase_pycall.jl")
+    end
 end
 
 @testset "Potential abstraction" begin
@@ -53,46 +76,54 @@ end
     @test potential(model, R) ≈ 2
 end
 
-@testset "AdiabaticModels" begin
-    @test test_model(Harmonic(), 10)
-    @test test_model(Free(), 10)
-    @test test_model(AveragedPotential((Harmonic(), Harmonic()), zeros(1, 10)), 10)
-    @test test_model(BosonBath(OhmicSpectralDensity(2.5, 0.1), 10), 10)
-    @test test_model(DarlingHollowayElbow(), 2)
-    @test test_model(Morse(), 1)
+if GROUP == "All" || GROUP == "Adiabatic"
+    @testset "AdiabaticModels" begin
+        @test test_model(Harmonic(), 10)
+        @test test_model(Free(), 10)
+        @test test_model(AveragedPotential((Harmonic(), Harmonic()), zeros(1, 10)), 10)
+        @test test_model(BosonBath(OhmicSpectralDensity(2.5, 0.1), 10), 10)
+        @test test_model(DarlingHollowayElbow(), 2)
+        @test test_model(Morse(), 1)
+    end
 end
 
-@testset "DiabaticModels" begin
-    @test test_model(DoubleWell(), 1)
-    @test test_model(TullyModelOne(), 1)
-    @test test_model(TullyModelTwo(), 1)
-    @test test_model(TullyModelThree(), 1)
-    @test test_model(Scattering1D(), 1)
-    @test test_model(ThreeStateMorse(), 1)
-    @test test_model(SpinBoson(DebyeSpectralDensity(0.25, 0.5), 10, 1.0, 1.0), 10)
-    @test test_model(OuyangModelOne(), 1)
-    @test test_model(GatesHollowayElbow(), 2)
-    @test test_model(MiaoSubotnik(Γ=0.1), 1)
-    @test test_model(AnanthModelOne(), 1)
-    @test test_model(AnanthModelTwo(), 1)
-    @test test_model(ErpenbeckThoss(Γ=2.0), 1)
-    @test test_model(WideBandBath(ErpenbeckThoss(Γ=2.0); step=0.1, bandmin=-1.0, bandmax=1.0), 1)
-    @test test_model(WideBandBath(GatesHollowayElbow(); step=0.1, bandmin=-1.0, bandmax=1.0), 2)
-    @test test_model(AndersonHolstein(ErpenbeckThoss(Γ=2.0), TrapezoidalRule(10, -1, 1)), 1)
-    @test test_model(AndersonHolstein(ErpenbeckThoss(Γ=2.0), ShenviGaussLegendre(10, -1, 1)), 1)
-    @test test_model(AndersonHolstein(GatesHollowayElbow(), ShenviGaussLegendre(10, -1, 1)), 2)
+if GROUP == "All" || GROUP == "Diabatic"
+    @testset "DiabaticModels" begin
+        @test test_model(DoubleWell(), 1)
+        @test test_model(TullyModelOne(), 1)
+        @test test_model(TullyModelTwo(), 1)
+        @test test_model(TullyModelThree(), 1)
+        @test test_model(Scattering1D(), 1)
+        @test test_model(ThreeStateMorse(), 1)
+        @test test_model(SpinBoson(DebyeSpectralDensity(0.25, 0.5), 10, 1.0, 1.0), 10)
+        @test test_model(OuyangModelOne(), 1)
+        @test test_model(GatesHollowayElbow(), 2)
+        @test test_model(MiaoSubotnik(Γ=0.1), 1)
+        @test test_model(AnanthModelOne(), 1)
+        @test test_model(AnanthModelTwo(), 1)
+        @test test_model(ErpenbeckThoss(Γ=2.0), 1)
+        @test test_model(WideBandBath(ErpenbeckThoss(Γ=2.0); step=0.1, bandmin=-1.0, bandmax=1.0), 1)
+        @test test_model(WideBandBath(GatesHollowayElbow(); step=0.1, bandmin=-1.0, bandmax=1.0), 2)
+        @test test_model(AndersonHolstein(ErpenbeckThoss(Γ=2.0), TrapezoidalRule(10, -1, 1)), 1)
+        @test test_model(AndersonHolstein(ErpenbeckThoss(Γ=2.0), ShenviGaussLegendre(10, -1, 1)), 1)
+        @test test_model(AndersonHolstein(GatesHollowayElbow(), ShenviGaussLegendre(10, -1, 1)), 2)
+    end
 end
 
-@testset "FrictionModels" begin
-    @test test_model(CompositeFrictionModel(Free(2), ConstantFriction(2, 1)), 3)
-    @test test_model(CompositeFrictionModel(Free(3), RandomFriction(3)), 3)
+if GROUP == "All" || GROUP == "Friction"
+    @testset "FrictionModels" begin
+        @test test_model(CompositeFrictionModel(Free(2), ConstantFriction(2, 1)), 3)
+        @test test_model(CompositeFrictionModel(Free(3), RandomFriction(3)), 3)
+    end
 end
 
-@testset "JuLIP" begin
-    using JuLIP: JuLIP
-    at = JuLIP.bulk(:Si, cubic=true)
-    deleteat!(at, 1)
-    JuLIP.set_calculator!(at, JuLIP.StillingerWeber())
-    model = AdiabaticModels.JuLIPModel(at)
-    @test test_model(model, length(at))
+if GROUP == "All" || GROUP == "JuLIP"
+    @testset "JuLIP" begin
+        using JuLIP: JuLIP
+        at = JuLIP.bulk(:Si, cubic=true)
+        deleteat!(at, 1)
+        JuLIP.set_calculator!(at, JuLIP.StillingerWeber())
+        model = AdiabaticModels.JuLIPModel(at)
+        @test test_model(model, length(at))
+    end
 end