fix

SciML · Oct 13, 2024 · 052dd9a · 052dd9a
1 parent 3274cb5
commit 052dd9a
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/src/discretize.jl b/src/discretize.jl
@@ -216,16 +216,17 @@ function generate_training_sets(domains, dx, eqs, bcs, eltypeθ, dict_indvars::D
     dict_var_span_ = Dict([Symbol(d.variables) => bc for (d, bc) in zip(domains, bc_data)])
 
     bcs_train_sets = map(bound_args) do bt
+        span = map(b -> get(dict_var_span, b, b), bt)
         _set = adapt(eltypeθ,
             hcat(vec(map(points -> collect(points), Iterators.product(span...)))...))
     end
 
     pde_vars = get_variables(eqs, dict_indvars, dict_depvars)
     pde_args = get_argument(eqs, dict_indvars, dict_depvars)
 
-    # pde_train_set = adapt(eltypeθ,
-    #     hcat(vec(map(points -> collect(points),
-    #         Iterators.product(bc_data...)))...))
+    pde_train_set = adapt(eltypeθ,
+        hcat(vec(map(points -> collect(points),
+            Iterators.product(bc_data...)))...))
 
     pde_train_sets = map(pde_args) do bt
         span = map(b -> get(dict_var_span_, b, b), bt)

diff --git a/src/neural_adapter.jl b/src/neural_adapter.jl
@@ -5,8 +5,8 @@ function generate_training_sets(domains, dx, eqs, eltypeθ)
         dxs = fill(dx, length(domains))
     end
     spans = [infimum(d.domain):dx:supremum(d.domain) for (d, dx) in zip(domains, dxs)]
-    _set = adapt(eltypeθ,
-        hcat(vec(map(points -> collect(points), Iterators.product(span...)))...))
+    train_set = adapt(eltypeθ,
+        hcat(vec(map(points -> collect(points), Iterators.product(spans...)))...))
 end
 
 function get_loss_function_(loss, init_params, pde_system, strategy::GridTraining)