From bed398bf858fc4083eb93dcb211486d5038f52fb Mon Sep 17 00:00:00 2001
From: Leandro Martinez <lmartine@unicamp.br>
Date: Mon, 7 Oct 2024 16:20:17 -0300
Subject: [PATCH] read and place fixed molecules

---
 src/data_structures/StructureType.jl | 127 ++++++++++++++++++++++++++-
 src/rigid_body/rigid_body.jl         |  29 ++++--
 test/data/diatomic.pdb               |   4 +
 3 files changed, 150 insertions(+), 10 deletions(-)
 create mode 100644 test/data/diatomic.pdb

diff --git a/src/data_structures/StructureType.jl b/src/data_structures/StructureType.jl
index 6347fe5..dfcb313 100644
--- a/src/data_structures/StructureType.jl
+++ b/src/data_structures/StructureType.jl
@@ -57,18 +57,21 @@ function read_structure_data(input_file_block::IOBuffer, tolerance;
         :atoms => nothing,
         :number_of_molecules => nothing,
         :reference_coordinates => nothing,
-        :constraints => Constraint[]
+        :constraints => Constraint[],
+        :fixed => (false, ),
+        :center => false,
     )
     # Read basic structure data first
     seekstart(input_file_block)
     atoms_block = false
     iconstraint = 0
+    seekstart(input_file_block)
     for line in eachline(input_file_block)
         keyword, values... = split(line)
         if keyword == "atoms"
             atoms_block = true
         end
-        if keyword == "end" && atoms_block
+        if atoms_block && keyword == "end"
             atoms_block = false
         end
         if keyword == "structure"
@@ -77,11 +80,15 @@ function read_structure_data(input_file_block::IOBuffer, tolerance;
             structure_data[:filename] = filename
             structure_data[:natoms] = length(atoms)
             structure_data[:atoms] = atoms
-            structure_data[:reference_coordinates] = set_reference_coordinates!(coor(atoms))
+            structure_data[:reference_coordinates] = coor(atoms)
             structure_data[:radii] = fill(T(tolerance), structure_data[:natoms])
             structure_data[:atom_constraints] = [Int[] for _ in 1:structure_data[:natoms]]
         elseif keyword == "number"
             structure_data[:number_of_molecules] = parse(Int, values[1]) 
+        elseif keyword == "fixed"
+            structure_data[:fixed] = (true, parse.(T, values))
+        elseif keyword == "center" 
+            structure_data[:center] = true
         elseif keyword in constraint_placements
             iconstraint += 1
             push!(structure_data[:constraints], parse_constraint["$keyword $(values[1])"](structure_data, values[2:end]; T=T))
@@ -94,6 +101,16 @@ function read_structure_data(input_file_block::IOBuffer, tolerance;
             structure_data[:radii] .= parse(T, values[1])
         end
     end
+    # If molecule is fixed, apply transformation to obtain the reference coordinates
+    if first(structure_data[:fixed])
+        position_fixed_molecule!(structure_data)
+        structure_data[:fixed] = true
+    else
+        if structure_data[:center]
+            throw(ArgumentError("option 'center' cannot be set without fixed position"))
+        end
+    end
+    pop!(structure_data, :center)
     #
     # Read custom atom radii and set constraints, according to atom blocks
     #
@@ -126,7 +143,24 @@ function read_structure_data(input_file_block::IOBuffer, tolerance;
     return StructureType{D,T}(;structure_data...)
 end
 
-@testitem "read_structure_data" begin
+# Functio to put fixed molecle in the position 
+# defined by the input file
+function position_fixed_molecule!(structure_data)
+    x = structure_data[:reference_coordinates]
+    cm = mean(x)
+    newcm = SVector(structure_data[:fixed][2][1:3]...)
+    rotation_angles = SVector(structure_data[:fixed][2][4:6]...)
+    move!(x, newcm, rotation_angles...)
+    # If the displacement is absolute, move the center of mass to the new position
+    if !structure_data[:center]
+        x .= x .+ Ref(cm)
+    end
+end
+
+@testitem "read_structure_data" setup=[RigidBody] begin
+    using Packmol
+    using PDBTools
+
     file = Packmol.src_dir*"/../test/data/water.pdb"
     tolerance = 2.0
 
@@ -183,6 +217,91 @@ end
     @test s.radii == Float32[1.0, 1.0, 2.0]
     @test s.constraints[1] == Box{Inside, Float32}([0.0, 0.0, 0.0], [40.0, 40.0, 40.0], 5.0)
     @test s.constraints[2] == Sphere{Outside, Float32}([0.0, 0.0, 0.0], 10.0, 5.0)
+
+
 end
 
+@testitem "fixed molecules" setup=[RigidBody] begin
+    using Packmol, PDBTools
+    file = Packmol.src_dir*"/../test/data/diatomic.pdb"
+    tolerance = 2.0
+
+    # Fixed molecule: do not move
+    input_file_block = """
+    structure $file        
+        number 1
+        fixed 0. 0. 0. 0. 0. 0.
+    end structure
+    """
+    s = Packmol.read_structure_data(input_file_block, tolerance) 
+    @test s.fixed == true
+    @test coor(s.atoms) ≈ s.reference_coordinates
+
+    # Fixed molecule without rotation: center of mass at origin
+    input_file_block = """
+    structure $file        
+        number 1
+        center
+        fixed 0. 0. 0. 0. 0. 0.
+    end structure
+    """
+    s = Packmol.read_structure_data(input_file_block, tolerance) 
+    @test s.reference_coordinates ≈ SVector{3, Float64}[[0.0, 0.0, -0.5], [0.0, 0.0, 0.5]] 
+
+    # Fixed molecule without rotation: center of mass at (10, 10, 10)
+    input_file_block = """
+    structure $file        
+        number 1
+        center
+        fixed 10. 10. 10. 0. 0. 0.
+    end structure
+    """
+    s = Packmol.read_structure_data(input_file_block, tolerance) 
+    @test s.reference_coordinates ≈ SVector{3, Float64}[[10.0, 10.0, 9.5], [10.0, 10.0, 10.5]]
+
+    # Rotate fixed molecule π/2 around x-axis (counterclockwise)
+    input_file_block = """
+    structure $file        
+        number 1
+        center
+        fixed 0. 0. 0. $(π/2) 0. 0.
+    end structure
+    """
+    s = Packmol.read_structure_data(input_file_block, tolerance) 
+    @test s.reference_coordinates ≈ SVector{3, Float64}[[0.0, 0.5, 0.0], [0.0, -0.5, 0.0]]
+
+    # Rotate fixed molecule π/2 around y-axis (counterclockwise)
+    input_file_block = """
+    structure $file        
+        number 1
+        center
+        fixed 0. 0. 0. 0. $(π/2) 0. 
+    end structure
+    """
+    s = Packmol.read_structure_data(input_file_block, tolerance) 
+    @test s.reference_coordinates ≈ SVector{3, Float64}[[-0.5, 0.0, 0.0], [0.5, 0.0, 0.0]]
+
+    # Rotate fixed molecule π/2 around z-axis (counterclockwise)
+    input_file_block = """
+    structure $file        
+        number 1
+        center
+        fixed 0. 0. 0. 0. 0. $(π/2)
+    end structure
+    """
+    s = Packmol.read_structure_data(input_file_block, tolerance) 
+    @test s.reference_coordinates ≈ SVector{3, Float64}[[0.0, 0.0, -0.5], [0.0, 0.0, 0.5]]
+
+    # Rotate fixed molecule π/2 around z-axis (counterclockwise) and move it to (10, 10, 10)
+    input_file_block = """
+    structure $file        
+        number 1
+        center
+        fixed 10. 10. 10. 0. 0. $(π/2)
+    end structure
+    """
+    s = Packmol.read_structure_data(input_file_block, tolerance) 
+    @test s.reference_coordinates ≈ SVector{3, Float64}[[10.0, 10.0, 9.5], [10.0, 10.0, 10.5]]
+
+end
 
diff --git a/src/rigid_body/rigid_body.jl b/src/rigid_body/rigid_body.jl
index 5ed1631..54e1c92 100644
--- a/src/rigid_body/rigid_body.jl
+++ b/src/rigid_body/rigid_body.jl
@@ -61,14 +61,21 @@ Translates and rotates a molecule according to the desired input center of coord
 =#
 function move!(x::AbstractVector{T}, newcm::T, beta, gamma, theta) where {T<:SVector}
     cm = mean(x)
+    x .= x .- Ref(cm)
+    rotate!(x, beta, gamma, theta)
+    x .= x .+ Ref(newcm)
+    return x
+end
+
+function rotate!(x::AbstractVector{T}, beta, gamma, theta) where {T<:SVector}
     A = eulermat(beta, gamma, theta)
     for i in eachindex(x)
-        x[i] = A * (x[i] - cm) + newcm
+        x[i] = A * x[i]
     end
     return x
 end
 
-@testitem "move!" setup=[RigidMolecules] begin
+@testitem "move!" setup=[RigidBody] begin
     x = [SVector(1.0, 0.0, 0.0), SVector(0.0, 0.0, 0.0)]
     @test Packmol.move!(copy(x), SVector(0.0, 0.0, 0.0), 0.0, 0.0, 0.0) ≈
           SVector{3,Float64}[[0.5, 0.0, 0.0], [-0.5, 0.0, 0.0]]
@@ -123,7 +130,7 @@ function random_move!(
 end
 
 
-@testitem "random_move!" setup=[RigidMolecules] begin
+@testitem "random_move!" setup=[RigidBody] begin
     # Orthorhombic cell
     x = [-1.0 .+ 2 * rand(SVector{3,Float64}) for _ = 1:5]
     system = ParticleSystem(positions=x, cutoff=0.1, unitcell=SVector(10.0, 10.0, 10.0), output=0.0)
@@ -148,10 +155,20 @@ end
 # of inertia (assuming all equal masses) along the z-axis. That will define the orientation
 # of the reference coordinates.
 #
-function set_reference_coordinates!(x::AbstractVector{<:SVector{3,T}}) where {T<:Real}
-    # Center of mass
+function set_reference_coordinates!(x::AbstractVector{<:SVector{3,T}}; fixed::Tuple{Bool, String, Vector{T}}) where {T<:Real}
+    # Move center of mass to the origin
     cm = mean(x)
     x .= x .- Ref(cm)
+    # If the molecule is fixed, apply the transformation that sets its position, and return
+    if first(fixed)
+        # Rotate molecule according to parameters of fixed position
+        rotate!(x, fixed[3][4], fixed[3][5], fixed[3][6])
+        if fixed[2] == "absolute"
+            x .= x + Ref(SVector(fixed[3][1:3]...))
+        end
+        return x
+    end
+    # If the molecule is not fixed, rotate it to align the principal moments of inertia
     # Inertia tensor
     I = zeros(MMatrix{3,3,T,9})
     for xi in x
@@ -199,7 +216,7 @@ end
     @test check_internal_distances(water, water_save)
 end
 
-@testsnippet RigidMolecules begin
+@testsnippet RigidBody begin
     using Packmol
     using CellListMap
     using StaticArrays
diff --git a/test/data/diatomic.pdb b/test/data/diatomic.pdb
new file mode 100644
index 0000000..f776787
--- /dev/null
+++ b/test/data/diatomic.pdb
@@ -0,0 +1,4 @@
+HEADER    diatomic
+HETATM    1  O   HOH     1       0.000   0.000   0.000
+HETATM    2  O   HOH     1       0.000   0.000   1.000
+END