Merge pull request #14 from mhauru/mhauru/fix-getindex-metricproblem

Make getindex on MetricProblem return a MetricProblem
Herb-AI · Oct 29, 2024 · f1e4a30 · f1e4a30
2 parents 4073698 + 0ab8b4b
commit f1e4a30
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/src/problem.jl b/src/problem.jl
@@ -99,6 +99,6 @@ end
 Overwrite `Base.getindex` to allow for slicing of input/output-based problems.
 """
 Base.getindex(p::Problem{Vector{IOExample}}, indices) = Problem(p.spec[indices])
-Base.getindex(p::MetricProblem{Vector{IOExample}}, indices) = Problem(p.spec[indices])
+Base.getindex(p::MetricProblem{Vector{IOExample}}, indices) = MetricProblem(p.cost_function, p.spec[indices])
 
 
diff --git a/test/test_ioproblem.jl b/test/test_ioproblem.jl
@@ -69,7 +69,8 @@ end
 
     # Test getindex
     submetric = metric2[1:2]
-    @test isa(submetric, Problem)
+    @test isa(submetric, MetricProblem)
     @test submetric.spec == spec[1:2]
     @test submetric.name == ""
+    @test submetric.cost_function === cost_function
 end