Refactor data type annotations in fit! and predict functions to use AbstractMatrix{<:Integer} instead of Matrix{Int}. Update LCAModel constructor to accept n_classes and n_items as Integer types for improved type flexibility. #4

Author: yanwenwang24

src/fit.jl

@@ -1,23 +1,23 @@
 """
-    fit!(model::LCAModel, data::Matrix{Int}; 
-         max_iter::Int=1000, tol::Float64=1e-6, verbose::Bool=false)
+    fit!(model::LCAModel, data::AbstractMatrix{<:Integer}; 
+         max_iter::Integer=1000, tol::Real=1e-6, verbose::Bool=false)
 
 Fit the LCA model using EM algorithm.
 
 # Arguments
 - `model::LCAModel`: Model to fit
-- `data::Matrix{Int}`: Prepared data matrix
-- `max_iter::Int=1000`: Maximum number of iterations
-- `tol::Float64=1e-6`: Convergence tolerance
+- `data::AbstractMatrix{<:Integer}`: Prepared data matrix
+- `max_iter::Integer=1000`: Maximum number of iterations
+- `tol::Real=1e-6`: Convergence tolerance
 - `verbose::Bool=false`: Whether to print progress
 
 # Returns
 - `Float64`: Final log-likelihood
 """
 
 function fit!(
-    model::LCAModel, data::Matrix{Int};
-    max_iter::Int=10000, tol::Float64=1e-6, verbose::Bool=false
+    model::LCAModel, data::AbstractMatrix{<:Integer};
+    max_iter::Integer=10000, tol::Real=1e-6, verbose::Bool=false
 )
     # Validate data dimensions
     n_obs, n_items = size(data)

src/predict.jl

@@ -1,19 +1,19 @@
 """
-    predict(model::LCAModel, data::Matrix{Int})
+    predict(model::LCAModel, data::AbstractMatrix{<:Integer})
 
 Predict class memberships for new data.
 
 # Arguments
 - `model::LCAModel`: Fitted model
-- `data::Matrix{Int}`: New data matrix
+- `data::AbstractMatrix{<:Integer}`: New data matrix
 
 # Returns
 - `Vector{Int}`: Predicted class assignments
 - `Matrix{Float64}`: Class membership probabilities
 """
 
 function predict(
-    model::LCAModel, data::Matrix{Int}
+    model::LCAModel, data::AbstractMatrix{<:Integer}
 )
     n_obs = size(data, 1)
     posterior = zeros(n_obs, model.n_classes)

src/types.jl

@@ -18,7 +18,7 @@ mutable struct LCAModel
     class_probs::Vector{Float64}
     item_probs::Vector{Matrix{Float64}}
 
-    function LCAModel(n_classes::Int, n_items::Int, n_categories::Vector{Int})
+    function LCAModel(n_classes::Integer, n_items::Integer, n_categories::AbstractVector{<:Integer})
         # Validate number of classes, items, and categories
         if n_classes < 2
             throw(ArgumentError("Number of classes must be ≥ 2, got $n_classes"))