diff --git a/Project.toml b/Project.toml
index 6752600..e94f1fd 100644
--- a/Project.toml
+++ b/Project.toml
@@ -4,22 +4,23 @@ authors = ["Br0kenSmi1e"]
 version = "1.0.0-DEV"
 
 [deps]
-HiGHS = "87dc4568-4c63-4d18-b0c0-bb2238e4078b"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+SCIP = "82193955-e24f-5292-bf16-6f2c5261a85f"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]
-HiGHS = "1"
 JuMP = "1"
 LinearAlgebra = "1"
+SCIP = "0.12"
 SpecialFunctions = "2"
 StaticArrays = "1"
 julia = "1"
 
 [extras]
+Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["Documenter", "Test"]
diff --git a/README.md b/README.md
index 5fe959f..acb92bf 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,5 @@
 # CrystalStructurePrediction
 
-[![Stable](https://img.shields.io/badge/docs-stable-blue.svg)](https://Br0kenSmi1e.github.io/CrystalStructurePrediction.jl/stable/)
-[![Dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://Br0kenSmi1e.github.io/CrystalStructurePrediction.jl/dev/)
 [![Build Status](https://github.com/Br0kenSmi1e/CrystalStructurePrediction.jl/actions/workflows/CI.yml/badge.svg?branch=main)](https://github.com/Br0kenSmi1e/CrystalStructurePrediction.jl/actions/workflows/CI.yml?query=branch%3Amain)
 [![Coverage](https://codecov.io/gh/Br0kenSmi1e/CrystalStructurePrediction.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/Br0kenSmi1e/CrystalStructurePrediction.jl)
 
@@ -27,5 +25,6 @@ If everything goes well, you should see the following crystal structure visualiz
 
 The center is Ti, surrounded by Sr (corner) and O (face center).
 
-## Reference
-[^Gusev2023] Gusev, V.V., Adamson, D., Deligkas, A., Antypov, D., Collins, C.M., Krysta, P., Potapov, I., Darling, G.R., Dyer, M.S., Spirakis, P., Rosseinsky, M.J., 2023. Optimality guarantees for crystal structure prediction. Nature 619, 68–72. https://doi.org/10.1038/s41586-023-06071-y
+## References
+
+[^Gusev2023]: Gusev, V.V., Adamson, D., Deligkas, A., Antypov, D., Collins, C.M., Krysta, P., Potapov, I., Darling, G.R., Dyer, M.S., Spirakis, P., Rosseinsky, M.J., 2023. Optimality guarantees for crystal structure prediction. Nature 619, 68–72. https://doi.org/10.1038/s41586-023-06071-y
diff --git a/examples/Project.toml b/examples/Project.toml
index 1e96f7d..8d23fa3 100644
--- a/examples/Project.toml
+++ b/examples/Project.toml
@@ -1,3 +1,4 @@
 [deps]
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 CrystalStructurePrediction = "4140a7b2-00d5-4ecf-8adb-c63b9db3f1fd"
+SCIP = "82193955-e24f-5292-bf16-6f2c5261a85f"
diff --git a/examples/SrTiO3-structure.png b/examples/SrTiO3-structure.png
index 32f9bab..e6df5a1 100644
Binary files a/examples/SrTiO3-structure.png and b/examples/SrTiO3-structure.png differ
diff --git a/examples/SrTiO3.jl b/examples/SrTiO3.jl
index 7ecd521..3c307c6 100644
--- a/examples/SrTiO3.jl
+++ b/examples/SrTiO3.jl
@@ -1,100 +1,41 @@
 using CrystalStructurePrediction
+using CairoMakie, SCIP
 
-"""
-    setup_crystal_parameters()
-
-Setup the parameters for SrTiO3 crystal structure prediction.
-Returns grid_size, population_list, species_list, charge_list, radii_list, lattice, depth, and alpha.
-"""
-function setup_crystal_parameters()
-    # Crystal structure parameters
-    grid_size = (2, 2, 2)
-    population_list = [1, 1, 3]  # 1 Sr, 1 Ti, 3 O atoms
-    species_list = [:Sr, :Ti, :O]
-    charge_list = [+2, +4, -2]
-    radii_list = [1.18, 0.42, 1.35]
-    
-    # Lattice parameters
-    lattice_constant = 3.899  # Å
-    L = lattice_constant * [1 0 0; 0 1 0; 0 0 1]
-    lattice = Lattice(L, (true, true, true))
-    
-    # Ewald summation parameters
-    depth = (4, 4, 4)
-    alpha = 2 / lattice_constant
-    
-    return grid_size, population_list, species_list, charge_list, radii_list, lattice, depth, alpha
-end
-
-"""
-    run_crystal_structure_prediction(; use_quadratic_problem::Bool=false)
-
-Run the crystal structure prediction for SrTiO3.
-
-# Arguments
-- `use_quadratic_problem::Bool`: Whether to use the quadratic problem. Only limited solvers support the quadratic problem, e.g. Gurobi (a commercial solver).
-"""
-function run_crystal_structure_prediction(; use_quadratic_problem::Bool=false)
-    # Setup parameters
-    grid_size, population_list, species_list, charge_list, radii_list, lattice, depth, alpha = setup_crystal_parameters()
-    
-    @info "Setting up crystal structure prediction for SrTiO3"
+# Run the crystal structure prediction, alpha is the Ewald parameter
+function run_crystal_structure_prediction(grid_size, populations, lattice, alpha; use_quadratic_problem::Bool=false)
+    @info "Setting up crystal structure prediction with"
     @info "Grid size: $grid_size"
-    @info "Population: $population_list $species_list"
-    @info "Charges: $charge_list"
-    @info "Ionic radii: $radii_list"
+    @info "Populations: $populations"
     
     # Build ion list and proximal pairs
-    ion_list = build_ion_list(grid_size, species_list, charge_list, radii_list)
+    ion_list = ions_on_grid(grid_size, collect(keys(populations)))
     @info "Created ion list with $(length(ion_list)) possible ion positions"
     
-    proximal_pairs = build_proximal_pairs(ion_list, lattice, 0.75)
-    @info "Identified $(length(proximal_pairs)) proximal pairs with cutoff 0.75"
-    
+    # Ewald summation parameters
+    depth = (4, 4, 4)
+    # Q: how to set alpha?
     if use_quadratic_problem
-        # Build interaction matrix
-        @info "Building interaction energy matrix..."
-        matrix = build_matrix(ion_list, lattice, interaction_energy, (alpha, depth, depth, depth))
-    
-        # Solve the quadratic problem
+        # Solve with the quadratic formulation
         @info "Solving quadratic optimization problem..."
-        energy, solution_x, csp = build_quadratic_problem(grid_size, population_list, matrix, proximal_pairs)
+        res = optimize_quadratic(ion_list, populations, lattice; optimizer=SCIP.Optimizer) do ion_a, ion_b, lattice
+            interaction_energy(ion_a, ion_b, lattice, alpha, depth, depth, depth)
+        end
     else
-        # Build interaction vector
-        @info "Building interaction energy vector..."
-        vector = build_vector(ion_list, lattice, interaction_energy, (alpha, depth, depth, depth))
-        
-        # Solve the linear problem
+        # Solve with the linear formulation
         @info "Solving linear optimization problem..."
-        energy, solution_x, csp = build_linear_problem(grid_size, population_list, vector, proximal_pairs)
-    end
-    
-    # Display results
-    @info "Optimization complete with energy: $energy"
-    @info "Solution structure:"
-    
-    selected_ions = []
-    for index in CartesianIndices(ion_list)
-        if solution_x[index] ≈ 1
-            push!(selected_ions, ion_list[index])
+        res = optimize_linear(ion_list, populations, lattice; optimizer=SCIP.Optimizer) do ion_a, ion_b, lattice
+            interaction_energy(ion_a, ion_b, lattice, alpha, depth, depth, depth)
         end
     end
-    
-    for (i, ion) in enumerate(selected_ions)
-        @info "Ion $i: $ion"
+    # Display results
+    @info "Optimization complete with energy: $(res.energy)"
+    for ion in res.selected_ions
+        @info "Ion: $ion"
     end
-    
-    return energy, selected_ions, csp
+    return res
 end
 
-# Run the prediction
-energy, selected_ions, csp = run_crystal_structure_prediction()
-
-# (-6.061349350569213, Any[Ion{3, Float64}(:Sr, 2, 1.18, [0.5, 0.5, 0.0]), Ion{3, Float64}(:Ti, 4, 0.42, [0.0, 0.0, 0.5]), Ion{3, Float64}(:O, -2, 1.35, [0.0, 0.0, 0.0]), Ion{3, Float64}(:O, -2, 1.35, [0.5, 0.0, 0.5]), Ion{3, Float64}(:O, -2, 1.35, [0.0, 0.5, 0.5])], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0  …  0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
-
 # Visualize the crystal structure
-using CairoMakie
-
 function visualize_crystal_structure(selected_ions, lattice, shift)
     fig = Figure(; size = (300, 250))
     ax = Axis3(fig[1, 1], 
@@ -128,11 +69,12 @@ function visualize_crystal_structure(selected_ions, lattice, shift)
     end
     
     # Plot the ions
-    properties = Dict(:Sr => (color = :green, size = 30), :Ti => (color = :aqua, size = 25), :O => (color = :red, size = 10))
+    properties = Dict(:Sr => (color = :green,), :Ti => (color = :aqua,), :O => (color = :red,))
     
     # Plot each ion
     for ion in selected_ions
         # Plot the ion at its position and all periodic images within the unit cell
+        coordinates = []
         for dx in -1:1, dy in -1:1, dz in -1:1
             # Add periodic image shift vector
             offset = [dx, dy, dz] .+ shift
@@ -141,18 +83,19 @@ function visualize_crystal_structure(selected_ions, lattice, shift)
             if all(0 .<= shifted_pos .<= 1)
                 # Convert to Cartesian coordinates
                 shifted_cart_pos = lattice.vectors * shifted_pos
-                scatter!(ax, [shifted_cart_pos[1]], [shifted_cart_pos[2]], [shifted_cart_pos[3]], 
-                         color = properties[ion.species].color, 
-                         markersize = properties[ion.species].size,
-                         label = string(ion.species))
+                push!(coordinates, shifted_cart_pos)
             end
         end
+        scatter!(ax, coordinates, 
+                    color = properties[ion.type.species].color, 
+                    markersize = ion.type.radii * 20,
+                    label = string(ion.type.species))
     end
  
     # Add legend with unique entries
-    unique_species = unique([ion.species for ion in selected_ions])
-    legend_elements = [MarkerElement(color = properties[sp].color, marker = :circle, markersize = properties[sp].size) for sp in unique_species]
-    legend_labels = [string(sp) for sp in unique_species]
+    unique_species = unique([ion.type for ion in selected_ions])
+    legend_elements = [MarkerElement(color = properties[sp.species].color, marker = :circle, markersize = sp.radii * 20) for sp in unique_species]
+    legend_labels = [string(sp.species) for sp in unique_species]
     
     Legend(fig[1, 2], legend_elements, legend_labels, "Species", patchsize = (30, 30))
     # Remove decorations and axis
@@ -161,10 +104,27 @@ function visualize_crystal_structure(selected_ions, lattice, shift)
     return fig
 end
 
-# Generate and save the visualization
-lattice = setup_crystal_parameters()[6]
-fig = visualize_crystal_structure(selected_ions, lattice, [0.5, 0.5, 0])
+####### Run the prediction #######
+function run_SrTiO3_prediction()
+    # Crystal structure parameters
+    lattice_constant = 3.899  # Å
+    lattice = Lattice(lattice_constant .* [1 0 0; 0 1 0; 0 0 1], (true, true, true))
+    grid_size = (2, 2, 2)
+    populations = Dict(
+        IonType(:Sr, +2, 1.18) => 1,  # 1 Sr atom
+        IonType(:Ti, +4, 0.42) => 1,  # 1 Ti atom
+        IonType(:O, -2, 1.35) => 3    # 3 O atoms
+    )
+
+    res = run_crystal_structure_prediction(grid_size, populations, lattice, 2.0/lattice_constant; use_quadratic_problem=false)
+
+    # Generate and save the visualization
+    origin = res.selected_ions[findfirst(x -> x.type.species == :Sr, res.selected_ions)].frac_pos
+    fig = visualize_crystal_structure(res.selected_ions, lattice, origin)
+
+    filename = joinpath(@__DIR__, "SrTiO3-structure.png")
+    save(filename, fig, dpi=20)
+    @info "Saved crystal structure visualization to: $filename"
+end
 
-filename = joinpath(@__DIR__, "SrTiO3-structure.png")
-save(filename, fig, dpi=20)
-@info "Saved crystal structure visualization to: $filename"
+run_SrTiO3_prediction()
diff --git a/src/CrystalStructurePrediction.jl b/src/CrystalStructurePrediction.jl
index eeed26f..bd90cce 100644
--- a/src/CrystalStructurePrediction.jl
+++ b/src/CrystalStructurePrediction.jl
@@ -3,17 +3,15 @@ module CrystalStructurePrediction
 using SpecialFunctions
 using LinearAlgebra
 using JuMP
-using HiGHS
+using SCIP
 using StaticArrays
 
-export Lattice, Ion
-export build_ion_list, build_vector
-export build_matrix, interaction_energy
-export build_linear_problem, build_quadratic_problem
+export Lattice, Ion, IonType
+export ions_on_grid, interaction_energy
+export optimize_linear, optimize_quadratic
 
 include("struct.jl")
 include("interaction.jl")
-include("build_matrix.jl")
 include("build_problem.jl")
 
 end
diff --git a/src/build_matrix.jl b/src/build_matrix.jl
deleted file mode 100644
index 67f26b9..0000000
--- a/src/build_matrix.jl
+++ /dev/null
@@ -1,53 +0,0 @@
-"""
-return a list of ions in the format of:
-`[ion(t1, p1), ion(t1, p2), ..., ion(t1,pn), ion(t2, p1), ..., ion(tm, pn)]`
-"""
-function build_ion_list(grid_size::NTuple{N, Int},
-        species_list::AbstractVector{Symbol},
-        charge_list::AbstractVector{Int},
-        radii_list::AbstractVector{Float64}
-        ) where N
-    return [Ion(species_list[t], charge_list[t], radii_list[t], (ci.I .- 1) ./ grid_size) for t in range(1, length(species_list)) for ci in vec(CartesianIndices(grid_size))]
-end
-
-function interaction_energy(
-        ion_a::Ion{D, T}, ion_b::Ion{D, T}, lattice::Lattice{D, T},
-        alpha::T, real_depth::NTuple{D, Int}, reciprocal_depth::NTuple{D, Int}, buckingham_depth::NTuple{D, Int}
-        ) where {D, T}
-    return real_space_sum(ion_a, ion_b, lattice, alpha, real_depth) + reciprocal_space_sum(ion_a, ion_b, lattice, alpha, reciprocal_depth) + buckingham_sum(ion_a, ion_b, lattice, buckingham_depth)# + radii_penalty(ion_a, ion_b, lattice, 0.75)
-end
-
-function build_matrix(
-        ion_list::AbstractVector{Ion{D, T}},
-        lattice::Lattice{D, T},
-        interaction_energy::FT,
-        parameters::Tuple
-        ) where {D, T<:Real, FT<:Function}
-    matrix = zeros(T, length(ion_list), length(ion_list))
-    for (index_a, ion_a) in enumerate(ion_list)
-        for (index_b, ion_b) in enumerate(ion_list[index_a:end])
-            matrix[index_a, index_a + index_b - 1] = interaction_energy(ion_a, ion_b, lattice, parameters...)
-        end
-    end
-    return (matrix + transpose(matrix)) / 2
-end
-
-function build_vector(
-        ion_list::AbstractVector{Ion{D, T}},
-        lattice::Lattice{D, T},
-        interaction_energy::FT,
-        parameters::Tuple
-        ) where {D, T<:Real, FT<:Function}
-    vector = zeros(T, length(ion_list) * (length(ion_list)-1) ÷ 2)
-    for (index_a, ion_a) in enumerate(ion_list)
-        for (index_b, ion_b) in enumerate(ion_list[index_a+1:end])
-            vector[index_a + (index_a+index_b-1)*(index_a+index_b-2)÷2] = interaction_energy(ion_a, ion_b, lattice, parameters...)
-        end
-    end
-    return vector
-end
-
-function build_proximal_pairs(ion_list::AbstractVector{Ion{D, T}}, lattice::Lattice{D, T}, c::Float64) where {D, T}
-    isProximal = (i, j) -> CrystalStructurePrediction.minimum_distance(ion_list[i], ion_list[j], lattice) < c * (ion_list[i].radii + ion_list[j].radii)
-    return [(i,j) for i in range(1, length(ion_list)) for j in range(i+1, length(ion_list)) if isProximal(i, j)]
-end
\ No newline at end of file
diff --git a/src/build_problem.jl b/src/build_problem.jl
index 711b6ff..f711e6e 100644
--- a/src/build_problem.jl
+++ b/src/build_problem.jl
@@ -1,100 +1,131 @@
-function build_problem(
-        grid_size::AbstractVector{Int},
-        population_list::AbstractVector{Int},
-        interaction_matrix::AbstractMatrix{T};
-        optimizer = HiGHS.Optimizer
-        ) where T<:Real
-    csp = Model(optimizer)
-    num_grid_points = prod(grid_size)
-    num_species = length(population_list)
-    @variable(csp, x[1:num_species, 1:num_grid_points], Bin)
-    @show interaction_matrix[48, 145]
-    @objective(csp, Min, sum(interaction_matrix[i,j] * x[(i-1)÷num_grid_points + 1, (i-1)%num_grid_points + 1] * x[(j-1)÷num_grid_points + 1, (j-1)%num_grid_points + 1] for i in range(1, num_grid_points*num_species) for j in range(1, num_grid_points*num_species)))
-    for i in range(1, num_species)
-        @constraint(csp, sum(x[i, j] for j in range(1, num_grid_points)) == population_list[i])
-    end
-    for j in range(1, num_grid_points)
-        @constraint(csp, sum(x[i, j] for i in range(1, num_species)) <= 1)
-    end
-    optimize!(csp)
-    assert_is_solved_and_feasible(csp)
-    return value.(x)
+struct IonOptimizationResult{D, T}
+    energy::T
+    selected_ions::AbstractVector{Ion{D, T}}
 end
 
-function build_linear_problem(
-        grid_size::NTuple{N, Int},
-        population_list::AbstractVector{Int},
-        interaction_vector::AbstractVector{T},
-        proximal_pairs::AbstractVector{Tuple{Int, Int}};
-        optimizer = HiGHS.Optimizer
-        ) where {N, T<:Real}
-    csp = Model(optimizer)
-    num_grid_points = prod(grid_size)
-    num_species = length(population_list)
-    @variable(csp, 0 <= x[1:num_species*num_grid_points] <= 1, Int)
-    @variable(csp, 0 <= s[1:num_species*num_grid_points*(num_species*num_grid_points-1)÷2] <= 1, Int)
+"""
+    optimize_linear(interaction, ions, populations, lattice; optimizer = SCIP.Optimizer, optimizer_options = Dict(), proximal_threshold = 0.75)
+
+Build a linear problem for crystal structure prediction. Suited for solvers not supporting quadratic constraints.
 
-    for t in range(1, num_species)
-        @constraint(csp, sum(x[num_grid_points*(t-1)+p] for p in range(1, num_grid_points)) == population_list[t])
-    end
-    for p in range(1, num_grid_points)
-        @constraint(csp, sum(x[num_grid_points*(t-1)+p] for t in range(1, num_species)) <= 1)
+# Arguments
+- `interaction::Function`: The interaction function.
+- `ions::AbstractVector{Ion{D, T}}`: The ions.
+- `populations::Dict{IonType{T}, Int}`: The populations of the ions.
+- `lattice::Lattice`: The lattice.
+
+# Keyword arguments
+- `optimizer`: The optimizer.
+- `optimizer_options`: The options for the optimizer, e.g. `optimizer_options = Dict("NodefileSave" => 1)` for Gurobi.
+- `proximal_threshold`: The threshold for the proximal pairs.
+"""
+function optimize_linear(interaction,
+            ions::AbstractVector{Ion{D, T}},
+            populations::Dict{IonType{T}, Int},  # a dictionary of ion types and their populations
+            lattice::Lattice;
+            optimizer = SCIP.Optimizer,
+            optimizer_options = Dict(),
+            proximal_threshold = 0.75
+        ) where {D, T<:Real}
+    # build the model
+    csp = Model(optimizer); set_silent(csp)
+    num_ions = length(ions)
+    unique_coordinates = unique!([ion.frac_pos for ion in ions])
+    @variable(csp, 0 <= x[1:num_ions] <= 1, Int)
+    @variable(csp, 0 <= s[1:num_ions*(num_ions-1)÷2] <= 1, Int)
+
+    # each ion type has a constraint on the number of ions of that type
+    for (type, population) in populations
+        @constraint(csp, sum(x[i] for i in range(1, num_ions) if ions[i].type == type) == population)
     end
-    for (i, j) in proximal_pairs
-        @constraint(csp, x[i] + x[j] <= 1)
+    # each grid point has at most one ion
+    for coord in unique_coordinates
+        @constraint(csp, sum(x[i] for i in range(1, num_ions) if ions[i].frac_pos == coord) <= 1)
     end
-    for i in range(1, num_species*num_grid_points-1)
-    for j in range(i+1, num_species*num_grid_points)
-        @constraint(csp, s[i + (j-1)*(j-2)÷2] <= x[i])
-        @constraint(csp, s[i + (j-1)*(j-2)÷2] <= x[j])
-        @constraint(csp, s[i + (j-1)*(j-2)÷2] >= x[i] + x[j] - 1)
+    # ions are not allowed to be too close to each other
+    for i in 1:num_ions-1, j in i+1:num_ions
+        if too_close(ions[i], ions[j], lattice, proximal_threshold)
+            @constraint(csp, x[i] + x[j] <= 1)
+        end
     end
+    # s is a "matrix" representing the co-existence of ions
+    energy = zero(eltype(s))
+    for i in range(1, num_ions-1)
+        for j in range(i+1, num_ions)
+            @constraint(csp, s[i + (j-1)*(j-2)÷2] <= x[i])
+            @constraint(csp, s[i + (j-1)*(j-2)÷2] <= x[j])
+            @constraint(csp, s[i + (j-1)*(j-2)÷2] >= x[i] + x[j] - 1)
+            energy += interaction(ions[i], ions[j], lattice) * s[i + (j-1)*(j-2)÷2]
+        end
     end 
-    @objective(csp, Min, dot(interaction_vector, s))
+    # minimize the interaction energy
+    @objective(csp, Min, energy)
+    # set the optimizer options
+    for (key, value) in optimizer_options
+        set_optimizer_attribute(csp, key, value)
+    end
+    # solve the problem
     optimize!(csp)
     assert_is_solved_and_feasible(csp)
-    return objective_value(csp), value.(x), value.(s)
+    return IonOptimizationResult(objective_value(csp), [ions[i] for i in 1:num_ions if value.(x)[i] ≈ 1])
+end
+
+function too_close(ion_a::Ion, ion_b::Ion, lattice::Lattice, c::Float64)
+    return minimum_distance(ion_a.frac_pos, ion_b.frac_pos, lattice) < c * (radii(ion_a) + radii(ion_b))
 end
 
 """
-    build_quadratic_problem(
-        grid_size::NTuple{N, Int},
-        population_list::AbstractVector{Int},
-        interaction_matrix::AbstractMatrix{T},
-        optimizer
-        ) where {N, T<:Real}
+    optimize_quadratic(interaction, ions, populations, lattice; optimizer = SCIP.Optimizer, optimizer_options = Dict(), proximal_threshold = 0.75)
 
 Build a quadratic problem for crystal structure prediction.
 
 # Arguments
-- `grid_size::NTuple{N, Int}`: The size of the grid.
-- `population_list::AbstractVector{Int}`: The number of atoms of each species.
-- `interaction_matrix::AbstractMatrix{T}`: The interaction matrix.
+- `interaction::Function`: The interaction function.
+- `ions::AbstractVector{Ion{D, T}}`: The ions.
+- `populations::Dict{IonType{T}, Int}`: The populations of the ions.
+- `lattice::Lattice`: The lattice.
+
+# Keyword arguments
 - `optimizer`: The optimizer.
+- `optimizer_options`: The options for the optimizer, e.g. `optimizer_options = Dict("NodefileSave" => 1)` for Gurobi.
+- `proximal_threshold`: The threshold for the proximal pairs.
 """
-function build_quadratic_problem(
-        grid_size::NTuple{N, Int},
-        population_list::AbstractVector{Int},
-        interaction_matrix::AbstractMatrix{T},
-        optimizer
-        ) where {N, T<:Real}
-
-    csp = Model(optimizer)
-    num_grid_points = prod(grid_size)
-    num_species = length(population_list)
-    @variable(csp, x[1:num_species*num_grid_points], Bin)
-    for t in range(1, num_species)
-        @constraint(csp, sum(x[num_grid_points*(t-1)+p] for p in range(1, num_grid_points)) == population_list[t])
+function optimize_quadratic(interaction,
+            ion_list::AbstractVector{Ion{D, T}},
+            populations::Dict{IonType{T}, Int},
+            lattice::Lattice;
+            optimizer = SCIP.Optimizer,
+            optimizer_options = Dict(),
+            proximal_threshold = 0.75
+        ) where {D, T<:Real}
+    # build the model
+    csp = Model(optimizer); set_silent(csp)
+    num_ions = length(ion_list)
+    unique_coordinates = unique!([ion.frac_pos for ion in ion_list])
+    @variable(csp, x[1:num_ions], Bin)
+    # each ion type has a constraint on the number of ions of that type
+    for (type, population) in populations
+        @constraint(csp, sum(x[i] for i in range(1, num_ions) if ion_list[i].type == type) == population)
+    end
+    # each grid point has at most one ion
+    for coord in unique_coordinates
+        @constraint(csp, sum(x[i] for i in range(1, num_ions) if ion_list[i].frac_pos == coord) <= 1)
+    end
+    # ions are not allowed to be too close to each other
+    for i in 1:num_ions-1, j in i+1:num_ions
+        if too_close(ion_list[i], ion_list[j], lattice, proximal_threshold)
+            @constraint(csp, x[i] + x[j] <= 1)
+        end
     end
-    @constraint(csp, sum(x[num_grid_points*(t-1)+1] for t in range(1, num_species)) == 1)
-    for p in range(2, num_grid_points)
-        @constraint(csp, sum(x[num_grid_points*(t-1)+p] for t in range(1, num_species)) <= 1)
+    # minimize the interaction energy
+    @objective(csp, Min, sum(interaction(ion_list[i], ion_list[j], lattice)*x[i]*x[j] for i in range(1, num_ions) for j in range(i+1, num_ions)))
+    # set the optimizer options
+    for (key, value) in optimizer_options
+        set_optimizer_attribute(csp, key, value)
     end
-    @objective(csp, Min, sum(interaction_matrix[i,j]*x[i]*x[j] for i in range(1, num_species*num_grid_points) for j in range(i+1, num_species*num_grid_points) if (j-i)%num_grid_points != 0))
-    set_optimizer_attribute(csp, "NodefileStart", 1)
-    
+    # solve the problem
     optimize!(csp)
     assert_is_solved_and_feasible(csp)
-    return objective_value(csp), value.(x), csp
+    return IonOptimizationResult(objective_value(csp), [ion_list[i] for i in 1:num_ions if value.(x)[i] ≈ 1])
 end
 
diff --git a/src/interaction.jl b/src/interaction.jl
index 5c729f0..415d713 100644
--- a/src/interaction.jl
+++ b/src/interaction.jl
@@ -1,79 +1,74 @@
-# a great short note:
-# https://www.cs.cornell.edu/courses/cs428/2006fa/Ewald%20Sum.pdf
-
-# ewald sum = real space sum + reciprocal space sum + "self" energy
-
-function build_shifts(depth::NTuple{N, Int}) where N
-    return [SVector((shift.I .- 1) .- depth) for shift in vec(CartesianIndices(2 .* depth .+ 1))]
-end
-
 """
-    summation(depth, interaction)
+    interaction_energy(ion_a, ion_b, lattice, alpha, real_depth, reciprocal_depth, buckingham_depth, buckingham_threshold, buckingham_penalty, radii_threshold, radii_penalty)
 
-Compute periodic summation of interaction to given depth.
+Compute the interaction energy between two ions, which is a sum of the real space Ewald sum, the reciprocal space Ewald sum, the Buckingham potential, and the radii penalty.
 
-Argument:
-    depth::AbstractVector{Int}: summation depth on the dimensions.
-    interaction<:Function: the interaction energy, given the shift as a list of integer.
-"""
-function summation(depth::NTuple{N, Int}, interaction::FT) where {N, FT<:Function}
-    return sum(interaction, build_shifts(depth))
-end
+# Arguments
+- `ion_a::Ion{D, T}`: the first ion.
+- `ion_b::Ion{D, T}`: the second ion.
+- `lattice::Lattice{D, T}`: the lattice.
+- `alpha::T`: the Ewald parameter.
+- `real_depth::NTuple{D, Int}`: summation depth on the dimensions for the real space Ewald sum.
+- `reciprocal_depth::NTuple{D, Int}`: summation depth on the dimensions for the reciprocal space Ewald sum.
+- `buckingham_depth::NTuple{D, Int}`: summation depth on the dimensions for the Buckingham potential.
 
-function real_space_potential(r::T, alpha::T) where T<:Real
-    return r ≈ 0 ? - alpha / (π)^0.5 : erfc(alpha * r) / (2 * r)
-end
+# Keyword arguments
+- `buckingham_threshold::T=0.75`: the threshold for the Buckingham potential.
+- `buckingham_penalty::T=3e2`: the penalty for the Buckingham potential.
+- `radii_threshold::T=0.0`: the threshold for the radii penalty, defined as the ratio of the distance between the two ions to the sum of their radii.
+- `radii_penalty::T=0.0`: the penalty for the radii penalty.
 
+# References
+- a great short note: https://www.cs.cornell.edu/courses/cs428/2006fa/Ewald%20Sum.pdf
 """
-    real_space_sum(depth, frac_pos_a, frac_pos_b, lattice, alpha)
-Compute the real space part of Ewald summation.
-
-Argument:
-    lattice::AbstractMatrix: in the form [a, b, c], where a, b, c are column vectors.
-"""
-function real_space_sum(
-        ion_a::Ion{D, T},
-        ion_b::Ion{D, T},
-        lattice::Lattice{D, T},
-        alpha::T,
-        depth::NTuple{D, Int}
-        ) where {D, T}
-    interaction = shift -> real_space_potential(norm(lattice.vectors * (ion_b.frac_pos + shift - ion_a.frac_pos)), alpha)
-    return ion_a.charge * ion_b.charge * 14.399645351950543 * summation(depth, interaction)
-end
-
-function reciprocal_space_potential(k::AbstractVector{T}, x::AbstractVector{T}, alpha::T) where T<:Real
-    return norm(k) ≈ 0 ? 0 : (2π / dot(k, k)) * exp(-( dot(k, k) / (4*alpha^2) ) ) * cos(dot(k, x))
-end
-
-function reciprocal_space_sum(
-        ion_a::Ion{D, T},
-        ion_b::Ion{D, T},
-        lattice::Lattice{D, T},
-        alpha::T,
-        depth::NTuple{D, Int}
+function interaction_energy(
+        ion_a::Ion{D, T}, ion_b::Ion{D, T}, lattice::Lattice{D, T},
+        alpha::T, real_depth::NTuple{D, Int}, reciprocal_depth::NTuple{D, Int}, buckingham_depth::NTuple{D, Int},
+        buckingham_threshold::T=0.75, buckingham_penalty::T=3e2, radii_threshold::T=0.0, radii_penalty::T=0.0
         ) where {D, T}
-    interaction = shift -> reciprocal_space_potential(2π * transpose(inv(lattice.vectors)) * shift, lattice.vectors * (ion_b.frac_pos - ion_a.frac_pos), alpha)
-    return ion_a.charge * ion_b.charge * 14.399645351950543 * summation(depth, interaction) / abs(det(lattice.vectors))
+    # Q: What is the constant 14.399645351950543?
+    # real space Ewald sum
+    energy = zero(T)
+    energy += charge(ion_a) * charge(ion_b) * 14.399645351950543 * periodic_sum(real_depth) do shift
+        real_space_potential(distance(lattice, ion_b.frac_pos + SVector(shift), ion_a.frac_pos), alpha)
+    end
+    # reciprocal space Ewald sum
+    energy += charge(ion_a) * charge(ion_b) * 14.399645351950543 / volume(lattice) * periodic_sum(reciprocal_depth) do shift
+        k = reciprocal_vectors(lattice) * SVector(shift)
+        norm(k) ≈ 0 && return zero(T)  # only sum over non-zero k
+        reciprocal_space_potential(k, cartesian(lattice, ion_b.frac_pos - ion_a.frac_pos), alpha)
+    end
+    # buckingham energy
+    energy += if ion_a == ion_b || minimum_distance(ion_a.frac_pos, ion_b.frac_pos, lattice) > buckingham_threshold * (radii(ion_a) + radii(ion_b))
+        periodic_sum(buckingham_depth) do shift
+            r = distance(lattice, ion_b.frac_pos + SVector(shift), ion_a.frac_pos)
+            r ≈ 0 && return zero(T)  # only sum over non-zero r
+            buckingham_potential(r, buckingham_parameters(ion_a.type, ion_b.type)...)
+        end
+    else
+        buckingham_penalty
+    end
+    # radii penalty
+    if !(ion_a == ion_b) && minimum_distance(ion_a.frac_pos, ion_b.frac_pos, lattice) / (radii(ion_a) + radii(ion_b)) > radii_threshold
+        energy += radii_penalty
+    end
+    return energy
 end
 
-function minimum_distance(ion_a::Ion{D, T}, ion_b::Ion{D, T}, lattice::Lattice{D, T}) where {D, T}
-    return min([norm(lattice.vectors * (ion_a.frac_pos - ion_b.frac_pos + shift)) for shift in build_shifts((1, 1, 1))]...)
+# What is this potential function? Does it have a name?
+function real_space_potential(r::T, alpha::T) where T<:Real
+    return T(r ≈ 0 ? -alpha / sqrt(π) : erfc(alpha * r) / (2 * r))
 end
 
-function radii_penalty(ion_a::Ion{D, T}, ion_b::Ion{D, T}, lattice::Lattice{D, T}, c::Float64, penalty::Float64=3e2) where {D, T}
-    if ion_a ≈ ion_b
-        return 0.0
-    else
-        return minimum_distance(ion_a, ion_b, lattice) / (ion_a.radii + ion_b.radii) > c ? 0.0 : penalty
-    end
+function reciprocal_space_potential(k::AbstractVector{T}, x::AbstractVector{T}, alpha::T) where T<:Real
+    return 2π / dot(k, k) * exp(-( dot(k, k) / (4*alpha^2) ) ) * cos(dot(k, x))
 end
 
 function buckingham_potential(r::T, A::T, ρ::T, C::T) where T<:Real
-    return r ≈ 0 ? 0 : A * exp(-r/ρ) - C / r^6
+    return A * exp(-r/ρ) - C / r^6
 end
 
-function buckingham_parameters(ion_a::Ion{D, T}, ion_b::Ion{D, T}) where {D, T}
+function buckingham_parameters(ion_a::IonType{T}, ion_b::IonType{T}) where T<:Real
     if Set([ion_a.species, ion_b.species]) == Set([:O])
         return 1388.7, 0.36262, 175.0
     elseif Set([ion_a.species, ion_b.species]) == Set([:Sr, :O])
@@ -81,24 +76,24 @@ function buckingham_parameters(ion_a::Ion{D, T}, ion_b::Ion{D, T}) where {D, T}
     elseif Set([ion_a.species, ion_b.species]) == Set([:Ti, :O])
         return 4590.7279, 0.261, 0.0
     else
+        # TODO: maybe throw an error here?
         return 0.0, 1.0, 0.0
     end
 end
 
-function buckingham_sum(
-        ion_a::Ion{D, T},
-        ion_b::Ion{D, T},
-        lattice::Lattice{D, T},
-        depth::NTuple{D, Int},
-        threshold::Float64=0.75,
-        penalty::Float64=3e2
-        ) where {D, T}
-    interaction = shift -> buckingham_potential(norm(lattice.vectors * (ion_b.frac_pos + shift - ion_a.frac_pos)), buckingham_parameters(ion_a, ion_b)...)
-    if ion_a ≈ ion_b
-        return summation(depth, interaction)
-    elseif minimum_distance(ion_a, ion_b, lattice) > threshold * (ion_a.radii + ion_b.radii)
-        return summation(depth, interaction)
-    else
-        return penalty
-    end
+"""
+    ions_on_grid(grid_size::NTuple{N, Int}, type_list::AbstractVector{IonType{T}}) where {N, T}
+
+Create a list of ions on a grid.
+
+# Arguments
+- `grid_size::NTuple{N, Int}`: the size of the grid.
+- `type_list::AbstractVector{IonType{T}}`: the list of ion types.
+
+# Returns
+- A vector of ions: `[ion(t1, p1), ion(t1, p2), ..., ion(t1,pn), ion(t2, p1), ..., ion(tm, pn)]`.
+where `t1, ..., tm` are the types of the ions and `p1, ..., pn` are the fractional positions of the ions on the grid.
+"""
+function ions_on_grid(grid_size::NTuple{N, Int}, type_list::AbstractVector{IonType{T}}) where {N, T}
+    return [Ion(type_list[t], SVector((ci.I .- 1) .// grid_size)) for t in range(1, length(type_list)) for ci in CartesianIndices(grid_size)]
 end
\ No newline at end of file
diff --git a/src/struct.jl b/src/struct.jl
index b89faa9..2c29fea 100644
--- a/src/struct.jl
+++ b/src/struct.jl
@@ -10,43 +10,90 @@ A lattice is a set of vectors that define the unit cell of a crystal.
 """
 struct Lattice{D, T, L}
     vectors::SMatrix{D, D, T, L}
+    # Q: why pbc is not used?
     pbc::NTuple{D, Bool}
 end
 function Lattice(vectors::AbstractMatrix{T}, pbc::NTuple{D, Bool}) where {D, T}
     return Lattice(SMatrix{D, D}(vectors), pbc)
 end
+# convert fractional coordinates to Cartesian coordinates
+cartesian(lt::Lattice, v) = lt.vectors * v
+# convert Cartesian coordinates to fractional coordinates
+fractional(lt::Lattice, v) = lt.vectors \ v
+# reciprocal lattice vectors
+reciprocal_vectors(lt::Lattice) = 2π .* transpose(inv(lt.vectors))
+# volume of the unit cell
+volume(lt::Lattice) = abs(det(lt.vectors))
+
+# the minimum distance between two ions in a lattice
+# TODO: this is the notoriously hard closest vector problem, try solve it with maybe integer programming?
+function minimum_distance(frac_pos_a::AbstractVector{Rational{Int}}, frac_pos_b::AbstractVector{Rational{Int}}, lattice::Lattice{D, T}) where {D, T}
+    return minimum(shift -> distance(lattice, frac_pos_b + SVector(shift), frac_pos_a), Iterators.product(ntuple(x->-1:1, D)...))
+end
+distance(a::AbstractVector{T}, b::AbstractVector{T}) where T = norm(a - b)
+distance(lt::Lattice{D, T}, frac_pos_a::AbstractVector{Rational{Int}}, frac_pos_b::AbstractVector{Rational{Int}}) where {D, T} = norm(cartesian(lt, frac_pos_b - frac_pos_a))
 
 """
-    Ion{D, T}
-    Ion(species::Symbol, charge::Int, radii::Float64, frac_pos)
+    periodic_sum(interaction, depth)
 
-An ion is a species with a charge, a radius, and a fractional position.
-`D` is the dimension of the space, and `T` is the type of the coordinates.
+Compute periodic summation of interaction to given depth.
+
+# Arguments
+- `interaction<:Function`: the interaction energy, given the shift as a list of integers.
+- `depth::NTuple{N, Int}`: summation depth on the dimensions.
+"""
+function periodic_sum(interaction::FT, depth::NTuple{N, Int}) where {N, FT<:Function}
+    return sum(interaction, Iterators.product(ntuple(i->-depth[i]:depth[i], N)...))
+end
+
+"""
+    IonType{T}
+    IonType(species::Symbol, charge::Int, radii::T)
+
+An ion type is a species with a charge and a radius.
+`T` is the type of the coordinates.
 
 # Fields
 - `species::Symbol`: The species of the ion.
 - `charge::Int`: The charge of the ion.
-- `radii::Float64`: The radius of the ion.
-- `frac_pos::SVector{D, T}`: The fractional position of the ion.
+- `radii::T`: The radius of the ion.
 
 # Examples
-To define an ion with species "O", charge -2, radius 1.35, and fractional position [0.25, 0.25, 0.25], run
-```julia
-julia> Ion(:O, -2, 1.35, [0.25, 0.25, 0.25])
-Ion{3, Float64}(:O, -2, 1.35, [0.25, 0.25, 0.25])
+To define an ion type with species "O", charge -2, radius 1.35, run
+```jldoctest
+julia> IonType(:O, -2, 1.35)
+IonType{Float64}(:O, -2, 1.35)
 ```
+
 """
-struct Ion{D, T}
+struct IonType{T}
     species::Symbol
     charge::Int
-    radii::Float64
-    frac_pos::SVector{D, T}
-end
-function Ion(species::Symbol, charge::Int, radii::Float64, frac_pos)
-    return Ion(species, charge, radii, SVector(frac_pos...))
+    radii::T
 end
+charge(ion::IonType) = ion.charge
+radii(ion::IonType) = ion.radii
+species(ion::IonType) = ion.species
+
+"""
+    Ion{D, T}
+    Ion(type::IonType{T}, frac_pos)
+
+An ion is a type with a fractional position.
+`D` is the dimension of the space, and `T` is the type of the coordinates.
 
-# filter out the periodic unit vectors
-periodic_vectors(lattice::Lattice) = lattice.vectors[:, findall(lattice.pbc)]
+# Fields
+- `type::IonType{T}`: The type of the ion.
+- `frac_pos::SVector{D, T}`: The fractional position of the ion.
+"""
+struct Ion{D, T}
+    type::IonType{T}
+    frac_pos::SVector{D, Rational{Int}}
+end
+function Ion(type::IonType{T}, frac_pos::AbstractVector{Rational{Int}}) where T
+    return Ion(type, SVector(frac_pos...))
+end
 
-Base.isapprox(a::Ion, b::Ion) = (a.species == b.species) && (a.charge == b.charge) && (a.radii ≈ b.radii) && (norm(a.frac_pos - b.frac_pos - floor.(a.frac_pos - b.frac_pos))≈0)
\ No newline at end of file
+charge(ion::Ion) = ion.type.charge
+radii(ion::Ion) = ion.type.radii
+species(ion::Ion) = ion.type.species
\ No newline at end of file
diff --git a/test/build_matrix.jl b/test/build_matrix.jl
deleted file mode 100644
index 80cd290..0000000
--- a/test/build_matrix.jl
+++ /dev/null
@@ -1,21 +0,0 @@
-using Test
-using LinearAlgebra: det
-using CrystalStructurePrediction
-using CrystalStructurePrediction: interaction_energy, build_matrix, build_ion_list
-
-@testset "interaction_energy" begin
-    lattice = Lattice(rand(3,3), (true, true, true))
-    ion_a = Ion(:O, -2, 1.35, rand(3))
-    ion_b = Ion(:O, -2, 1.35, rand(3))
-    alpha = 2.0 / (abs(det(lattice.vectors)))^(1/3)
-    depth = (4, 4, 4)
-    @test interaction_energy(ion_a, ion_b, lattice, alpha, depth, depth, depth) ≈ interaction_energy(ion_b, ion_a, lattice, alpha, depth, depth, depth)
-end
-
-@testset "build_matrix" begin
-    lattice = Lattice(rand(2, 2), (true, true))
-    ion_list = build_ion_list((1, 1), [:None], [1], [0.0])
-    alpha = 2.0 / (abs(det(lattice.vectors)))^(1/2)
-    depth = (0, 0)
-    @test build_matrix(ion_list, lattice, interaction_energy, (alpha, depth, depth, depth)) ≈ [-14.399645351950543*alpha/(π)^0.5]
-end
\ No newline at end of file
diff --git a/test/build_problem.jl b/test/build_problem.jl
new file mode 100644
index 0000000..c6212f6
--- /dev/null
+++ b/test/build_problem.jl
@@ -0,0 +1,33 @@
+using CrystalStructurePrediction, Test
+
+@testset "build_problem" begin
+    # Crystal structure parameters
+    grid_size = (2, 2, 2)
+    population_list = [1, 1, 3]  # 1 Sr, 1 Ti, 3 O atoms
+    Sr, Ti, O = IonType(:Sr, 2, 1.18), IonType(:Ti, 4, 0.42), IonType(:O, -2, 1.35)
+    type_list = [Sr, Ti, O]
+    
+    # Lattice parameters
+    lattice_constant = 3.899  # Å
+    L = lattice_constant * [1 0 0; 0 1 0; 0 0 1]
+    lattice = Lattice(L, (true, true, true))
+    
+    # Ewald summation parameters
+    depth = (4, 4, 4)
+    alpha = 2 / lattice_constant
+    
+    # Build ion list and proximal pairs
+    ion_list = ions_on_grid(grid_size, type_list)
+    
+    # Solve the linear problem
+    res_linear = optimize_linear(ion_list, Dict(type_list .=> population_list), lattice) do ion_a, ion_b, lattice
+        interaction_energy(ion_a, ion_b, lattice, alpha, depth, depth, depth)
+    end
+    @test res_linear.energy ≈ -6.061349350569214
+
+    # The quadratic problem formulation
+    res_quadratic = optimize_quadratic(ion_list, Dict(type_list .=> population_list), lattice) do ion_a, ion_b, lattice
+        interaction_energy(ion_a, ion_b, lattice, alpha, depth, depth, depth)
+    end
+    @test res_quadratic.energy ≈ -6.061349350569214
+end
\ No newline at end of file
diff --git a/test/interaction.jl b/test/interaction.jl
index 6cedf8d..f639bc6 100644
--- a/test/interaction.jl
+++ b/test/interaction.jl
@@ -1,40 +1,19 @@
 using Test
+using LinearAlgebra: det
 using CrystalStructurePrediction, StaticArrays
-using CrystalStructurePrediction: interaction_energy, real_space_sum, reciprocal_space_sum, buckingham_sum
+using CrystalStructurePrediction: interaction_energy, ions_on_grid
 
 @testset "lattice" begin
     lattice = Lattice(rand(3,3), (true, true, true))
-    @test lattice.vectors ≈ CrystalStructurePrediction.periodic_vectors(lattice)
+    @test lattice isa Lattice
 end
 
-@testset "real_space_sum" begin
-    ion_a = Ion(:none, 1, 0.0, rand(3))
-    ion_b = Ion(:none, 2, 0.0, rand(3))
-    lattice = Lattice(rand(3,3), (true, true, true))
-    alpha = 2.0
-    for depth in range(0,4)
-        depth_list = ntuple(x->depth, 3)
-        @test real_space_sum(ion_a, ion_b, lattice, alpha, depth_list) ≈ real_space_sum(ion_b, ion_a, lattice, alpha, depth_list)
-    end
-end
-
-@testset "reciprocal_space_sum" begin
-    ion_a = Ion(:none, 1, 0.0, rand(3))
-    ion_b = Ion(:none, 2, 0.0, rand(3))
-    lattice = Lattice(rand(3,3), (true, true, true))
-    alpha = 2.0
-    for depth in range(0,4)
-        depth_list = ntuple(x->depth, 3)
-        @test reciprocal_space_sum(ion_a, ion_b, lattice, alpha, depth_list) ≈ reciprocal_space_sum(ion_b, ion_a, lattice, alpha, depth_list)
-    end
-end
-
-@testset "buckingham_sum" begin
-    ion_a = Ion(:Sr, 1, 0.0, rand(3))
-    ion_b = Ion(:O, 2, 0.0, rand(3))
-    lattice = Lattice(rand(3,3), (true, true, true))
-    for depth in range(0,4)
-        depth_list = ntuple(x->depth, 3)
-        @test buckingham_sum(ion_a, ion_b, lattice, depth_list) ≈ buckingham_sum(ion_b, ion_a, lattice, depth_list)
-    end
+@testset "interaction_energy" begin
+    lattice = Lattice([1.0 0.0 0.0; 0.0 0.8 0.0; 0.0 0.6 1.2], (true, true, true))
+    ion_a = Ion(IonType(:O, -2, 1.35), [1, 1//2, 1//2])
+    ion_b = Ion(IonType(:O, -2, 1.35), [0, 0, 1//5])
+    alpha = 2.0 / (abs(det(lattice.vectors)))^(1/3)
+    depth = (4, 4, 4)
+    @test interaction_energy(ion_a, ion_b, lattice, alpha, depth, depth, depth) ≈ interaction_energy(ion_b, ion_a, lattice, alpha, depth, depth, depth)
+    @test interaction_energy(ion_a, ion_b, lattice, alpha, depth, depth, depth) ≈ 322.4479207721009
 end
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index 146ddf6..f0f5424 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,10 +1,18 @@
 using CrystalStructurePrediction
-using Test
+using Test, Documenter
+
+@testset "struct.jl" begin
+    include("struct.jl")
+end
 
 @testset "interaction.jl" begin
     include("interaction.jl")
 end
 
-@testset "build_matrix.jl" begin
-    include("build_matrix.jl")
+@testset "build_problem.jl" begin
+    include("build_problem.jl")
 end
+
+# Run doctests
+DocMeta.setdocmeta!(CrystalStructurePrediction, :DocTestSetup, :(using CrystalStructurePrediction); recursive=true)
+doctest(CrystalStructurePrediction)
\ No newline at end of file
diff --git a/test/struct.jl b/test/struct.jl
new file mode 100644
index 0000000..b5d6158
--- /dev/null
+++ b/test/struct.jl
@@ -0,0 +1,34 @@
+using CrystalStructurePrediction, Test
+using CrystalStructurePrediction: species, charge, radii, fractional, cartesian, minimum_distance, distance
+
+@testset "struct.jl" begin
+    # Test IonType constructor
+    ion_type = IonType(:O, -2, 1.35)
+    @test species(ion_type) == :O
+    @test charge(ion_type) == -2
+    @test radii(ion_type) == 1.35
+    
+    # Test Ion constructor
+    ion = Ion(ion_type, [1//2, 1//2, 1//2])
+    @test species(ion) == :O
+    @test charge(ion) == -2
+    @test radii(ion) == 1.35
+    @test ion.frac_pos ≈ [1//2, 1//2, 1//2]
+    
+    # Test Lattice constructor
+    vectors = [0.5 0 0; 0 0.5 0; 0 0 0.5]
+    pbc = (true, true, true)
+    lattice = Lattice(vectors, pbc)
+    @test lattice.vectors ≈ vectors
+    @test lattice.pbc == pbc
+    @test minimum_distance([0//1, 0, 0], [9//10, 9//10, 9//10], lattice) ≈ 0.05 * sqrt(3)
+    @test distance([0//1, 0, 0], [9//10, 9//10, 9//10]) ≈ sqrt(3) * 0.9
+    @test distance(lattice, [0//1, 0, 0], [9//10, 9//10, 9//10]) ≈ sqrt(3) * 0.45
+
+    # Test fractional and cartesian conversion
+    vectors = randn(3, 3)
+    lattice = Lattice(vectors, pbc)
+    frac_pos = rand(3)
+    cart_pos = cartesian(lattice, frac_pos)
+    @test fractional(lattice, cart_pos) ≈ frac_pos
+end
\ No newline at end of file