Changed module to actually function and started adding testing.

julia_code/.vscode/extensions.json

@@ -0,0 +1,13 @@
+{
+	// See https://go.microsoft.com/fwlink/?LinkId=827846 to learn about workspace recommendations.
+	// Extension identifier format: ${publisher}.${name}. Example: vscode.csharp
+
+	// List of extensions which should be recommended for users of this workspace.
+	"recommendations": [
+		
+	],
+	// List of extensions recommended by VS Code that should not be recommended for users of this workspace.
+	"unwantedRecommendations": [
+		
+	]
+}

julia_code/.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "python.pythonPath": "/bin/python3"
+}

julia_code/Module/notebooks/ExampleUse.jl

@@ -0,0 +1,81 @@
+include("../src/Orbits.jl")
+using .Orbits
+using LinearAlgebra
+
+#Set key dimensions
+const N_constellations = 4
+const N_debris = 1
+const N_states = N_constellations + N_debris
+
+#Setup Economic Models
+em2_a = LinearModel(0.95, [1 -0.02 -0.02 0], [5.0 0 0 0])
+em2_b = LinearModel(0.95, [-0.02 1 -0.02 0], [0.0 5 0 0])
+em2_c = LinearModel(0.95, [0 -0.02 1 -0.02], [0.0 0 5 0])
+em2_d = LinearModel(0.95, [0 -0.02 -0.02 1], [0.0 0 0 5])
+
+#Setup Physics
+
+basic_model = BasicPhysics(
+    0.002
+    ,0.002*(ones(N_constellations,N_constellations) - LinearAlgebra.I)
+    ,0.01
+    ,0.001
+    ,5.0
+    ,0.05
+)
+
+# Setup NN
+bg = BranchGenerator(N_constellations)
+operators_policy = bg(operator_policy_function_generator(N_constellations,N_debris),vcat)
+planners_policy = bg(operator_policy_function_generator(N_constellations,N_debris),vcat)
+planners_value = value_function_generator(64)
+
+
+# Setup Operators
+const operator_array = [ 
+    #first operator
+    OperatorLoss(		
+        em2_a
+        ,value_function_generator()
+        ,operators_policy #this is held by all operators
+        ,params(operators_policy[2][1]) #first operator gets first branch of params
+        ,basic_model
+    )
+    ,OperatorLoss(		
+        em2_b
+        ,value_function_generator()
+        ,operators_policy #this is held by all operators
+        ,params(operators_policy[2][2]) #first operator gets first branch of params
+        ,basic_model
+    )
+    ,OperatorLoss(		
+        em2_c
+        ,value_function_generator()
+        ,operators_policy #this is held by all operators
+        ,params(operators_policy[2][3]) #first operator gets first branch of params
+        ,basic_model
+    )
+    ,OperatorLoss(		
+        em2_d
+        ,value_function_generator()
+        ,operators_policy #this is held by all operators
+        ,params(operators_policy[2][4]) #first operator gets first branch of params
+        ,basic_model
+    )
+]
+
+#sanity check time
+@assert length(operator_array) == N_constellations "Mismatch in predetermined number of constellations and the number of operators initialized"
+
+# Setup Planner
+pl = PlannerLoss(
+		0.95
+		,operator_array
+		,planners_policy
+		,params(planners_policy)
+		,planners_value
+		,params(planners_value)
+        ,basic_model
+)
+
+# Export Planner

julia_code/Module/src/Constellations.jl

@@ -0,0 +1,401 @@
+
+module Orbits
+
+
+#Add exports here
+export State ,state_to_tuple ,state_transition ,PhysicalModel ,BasicPhysics
+
+export BranchGenerator, value_function_generator, 
+    cross_linked_planner_policy_function_generator,
+    operator_policy_function_generator
+
+# Exports
+export GeneralizedLoss ,OperatorLoss ,PlannerLoss ,UniformDataConstructor,
+        LinearModel ,BasicPhysics,
+        operator_policy_function_generator ,value_function_generator ,BranchGenerator,
+        TupleDuplicator, BranchGenerator
+
+using Flux, LinearAlgebra
+using BSON, Zygote
+export LinearModel, BenefitFunction, EconomicModel
+
+#==#
+struct State
+        stocks::Array{Float32}
+        debris::Array{Float32}
+end
+
+function state_to_tuple(s::State)
+        return (s.stocks ,s.debris)
+end
+
+
+
+abstract type PhysicalModel end
+        
+struct BasicPhysics <: PhysicalModel
+        #rate at which debris hits satellites
+        debris_collision_rate::Real
+        #rate at which satellites of different constellations collide
+        satellite_collision_rates::Matrix{Float64}
+        #rate at which debris exits orbits
+        decay_rate::Real
+        #rate at which satellites
+        autocatalysis_rate::Real
+        #ratio at which a collision between satellites produced debris
+        satellite_collision_debris_ratio::Real
+        #Ratio at which launches produce debris
+        launch_debris_ratio::Real
+end
+function state_transition(
+        physics::BasicPhysics 
+        ,state::State 
+        ,launches::Vector{Float32}
+        )
+        #=
+        Physical Transitions
+        =#
+        survival_rate = survival(state,physics)
+
+        # Debris transitions
+
+        # get changes in debris from natural dynamics
+        natural_debris_dynamics = (1 - physics.decay_rate + physics.autocatalysis_rate) * state.debris 
+        
+        # get changes in debris from satellite loss
+        satellite_loss_debris = physics.satellite_collision_debris_ratio * (1 .- survival_rate)' * state.stocks 
+        
+        # get changes in debris from launches
+        launch_debris = physics.launch_debris_ratio * sum(launches)
+        
+        # total debris level
+        debris′ = natural_debris_dynamics .+ satellite_loss_debris .+ launch_debris
+
+
+        # Stocks Transitions
+        stocks′ = (LinearAlgebra.diagm(survival_rate) .- physics.decay_rate)*state.stocks + launches
+
+
+        return State(stocks′,debris′)
+end 
+
+
+function survival(
+                state::State
+                ,physical_model::BasicPhysics
+                )    
+        return exp.(-(physical_model.satellite_collision_rates .+ physical_model.decay_rate) * state.stocks .- (physical_model.debris_collision_rate * state.debris))
+end
+
+
+
+
+#=
+Benefit Functions:
+
+These represent the benefits that different operators may recieve
+=#
+abstract type EconomicModel end
+
+#basic linear model
+struct LinearModel <: EconomicModel
+        β::Float32
+        payoff_array::Array{Float32}
+        policy_costs::Array{Float32}
+end
+function (em::LinearModel)(
+        state::State
+        ,actions::Vector{Float32}
+        )
+        return em.payoff_array*state.stocks - em.policy_costs*actions
+end
+
+#basic CES model
+struct CES <: EconomicModel
+        β::Float32
+        r::Float32 #elasticity of subsititution
+        payoff_array::Array{Float32}
+        policy_costs::Array{Float32}
+        debris_costs::Array{Float32}
+end
+function (em::CES)(
+        state::State
+        ,actions::Vector{Float32}
+        )
+
+        #issue here with multiplication
+        r1 = em.payoff_array .* (state.stocks.^em.r)
+        r2 = - em.debris_costs .* (state.debris.^em.r)
+        r3 = - em.policy_costs .* (actions.^em.r)
+        return (r1 + r2 + r3) .^ (1/em.r)
+
+end
+
+#basic CRRA
+struct CRRA <: EconomicModel
+        β::Float32
+        σ::Float32 #elasticity of subsititution
+        payoff_array::Array{Float32}
+        policy_costs::Array{Float32}
+        debris_costs::Array{Float32}
+end
+function (em::CRRA)(
+        state::State
+        ,actions::Vector{Float32}
+        )
+        #issue here with multiplication
+        core = (em.payoff_array*state.stocks - em.debris_costs*state.debris - em.policy_costs*actions).^(1 - em.σ)
+
+        return (core-1)/(1-em.σ)
+end
+
+
+
+#= TupleDuplicator
+This is used to create a tuple of size n with deepcopies of any object x
+=#
+struct TupleDuplicator
+        n::Int
+end
+#make tuples consisting of copies of whatever was provided
+function (f::TupleDuplicator)(s::State)
+    st = state_to_tuple(s)
+    return f(st)
+    #BROKEN: Fails in test, but works in REPL.
+end
+
+(f::TupleDuplicator)(x) = tuple([deepcopy(x) for i=1:f.n]...)
+
+
+
+struct BranchGenerator
+        n::UInt
+end
+function (b::BranchGenerator)(branch::Flux.Parallel,join_fn::Function) 
+        # used to deepcopy the branches and duplicate the inputs in the returned chain
+        f = TupleDuplicator(b.n)
+        
+        return Flux.Chain(state_to_tuple,f,Flux.Parallel(join_fn,f(branch)...))
+        #note that it destructures the state to a tuple, duplicates,
+        # and then passes to the parallelized functions.
+end
+
+
+# Neural Network Generators
+
+function value_function_generator(number_params=32)
+    return Flux.Chain(
+        Flux.Parallel(vcat
+                #parallel joins together stocks and debris, after a little bit of preprocessing
+                ,Flux.Chain(
+                    Flux.Dense(N_constellations, number_params*2,Flux.relu)
+                    ,Flux.Dense(number_params*2, number_params*2,Flux.σ)
+                )
+                ,Flux.Chain(
+                    Flux.Dense(N_debris, number_params,Flux.relu)
+                    ,Flux.Dense(number_params, number_params,Flux.σ)
+                )
+            )
+            #Apply some transformations to the preprocessed data.
+            ,Flux.Dense(number_params*3,number_params,Flux.σ)
+            ,Flux.Dense(number_params,number_params,Flux.σ)
+            ,Flux.Dense(number_params,1)
+        )
+
+end
+
+function cross_linked_planner_policy_function_generator(number_params=32)
+    return Flux.Chain(
+        Flux.Parallel(vcat
+                #parallel joins together stocks and debris
+                ,Flux.Chain(
+                    Flux.Dense(N_constellations, number_params*2,Flux.relu)
+                    ,Flux.Dense(number_params*2, number_params*2,Flux.σ)
+                )
+                ,Flux.Chain(
+                    Flux.Dense(N_debris, number_params,Flux.relu)
+                    ,Flux.Dense(number_params, number_params)
+                )
+            )
+            #Apply some transformations
+            ,Flux.Dense(number_params*3,number_params,Flux.σ)
+            ,Flux.Dense(number_params,number_params,Flux.σ)
+            ,Flux.Dense(number_params,N_constellations,Flux.relu)
+        )
+
+end
+
+function operator_policy_function_generator(
+    N_constellations::Int
+    ,N_debris
+    ,number_params=32
+    )
+    return Flux.Chain(
+        Flux.Parallel(vcat
+                #parallel joins together stocks and debris
+                ,Flux.Chain(
+                    Flux.Dense(N_constellations, number_params*2,Flux.relu)
+                    ,Flux.Dense(number_params*2, number_params*2,Flux.tanh)
+                )
+                ,Flux.Chain(
+                    Flux.Dense(N_debris, number_params,Flux.relu)
+                    ,Flux.Dense(number_params, number_params)
+                )
+            )
+            #Apply some transformations
+            ,Flux.Dense(number_params*3,number_params,Flux.tanh)
+            ,Flux.Dense(number_params,number_params,Flux.tanh)
+            ,Flux.Dense(number_params,1,x -> Flux.relu.(sinh.(x)))
+        )
+
+end
+
+
+
+
+
+abstract type GeneralizedLoss end
+
+#=
+This struct organizes information about a given constellation operator
+Used to provide an interable loss function for training
+=#
+struct OperatorLoss <: GeneralizedLoss
+        #econ model describing operator
+        econ_model::EconomicModel
+        #Operator's value and policy functions, as well as which parameters the operator can train.
+        operator_value_fn::Flux.Chain
+        collected_policies::Flux.Chain #this is held by all operators
+        operator_policy_params::Flux.Params
+        physics::PhysicalModel
+end
+# overriding function to calculate operator loss
+function (operator::OperatorLoss)(
+        state::State
+)
+
+
+        #get actions
+        a = operator.collected_policies((s,d))
+
+        #get updated stocks and debris
+        state′ = state_transition(operator.physics,state,a)
+
+
+        bellman_residuals = operator.operator_value_fn(state) - operator.econ_model(state,a) - operator.econ_model.β * operator.operator_value_fn(state′)
+
+        maximization_condition = - operator.econ_model(state ,a) - operator.econ_model.β * operator.operator_value_fn((state′))
+
+        return Flux.mae(bellman_residuals.^2 ,maximization_condition)
+end # struct
+function (operator::OperatorLoss)(
+        state::State
+        ,policy::Flux.Chain #allow for another policy to be subsituted
+        )
+
+
+        #get actions
+        a = policy(state)
+
+        #get updated stocks and debris
+        state′ = stake_transition(state,a)
+
+
+        bellman_residuals = operator.operator_value_fn(state) - operator.econ_model(state,a) - operator.econ_model.β * operator.operator_value_fn(state′)
+
+        maximization_condition = - operator.econ_model(state,a) - operator.econ_model.β * operator.operator_value_fn(state′)
+
+        return Flux.mae(bellman_residuals.^2 ,maximization_condition)
+end #function
+
+function train_operators!(op::OperatorLoss, data::State, opt)
+        Flux.train!(op, op.operator_policy_params, data, opt)
+        Flux.train!(op, Flux.params(op.operator_value_fn), data, opt)
+end
+
+
+#=
+Describe the Planners loss function
+=#
+struct PlannerLoss <: GeneralizedLoss
+        #=
+        Ideally, with just a well formed PlannerLoss and the training functions below, we should be able to train the approximation.
+
+        There is an issue with appropriately training the value functions.
+        In this case, it is not happening...
+        =#
+        β::Float32
+
+        operators::Array{GeneralizedLoss}
+
+        policy::Flux.Chain
+        policy_params::Flux.Params
+        value::Flux.Chain
+        value_params::Flux.Params
+
+        physical_model::PhysicalModel
+end
+function (planner::PlannerLoss)(
+        state::State
+)
+        a = planner.policy((s ,d))
+        #get updated stocks and debris
+        state′ = state_transition(s ,d ,a)
+
+
+
+                #calculate the total benefit from each of the models
+        benefit = sum([ co.econ_model(s ,d ,a) for co in planner.operators])
+                #issue here with mutating. Maybe generators/list comprehensions?
+
+
+        bellman_residuals = planner.value((s,d)) - benefit - planner.β .* planner.value((s′,d′))
+
+        maximization_condition = - benefit - planner.β .* planner.value((s′,d′))
+
+        return Flux.mae(bellman_residuals.^2 ,maximization_condition)
+end
+
+function train_planner!(pl::PlannerLoss, N_epoch::Int, opt)
+        errors = []
+        for i = 1:N_epoch
+            data = data_gen()
+            Flux.train!(pl, pl.policy_params, data, opt)
+            Flux.train!(pl, pl.value_params, data, opt)
+            append!(errors, error(pl, data1) / 200)
+        end
+        return errors
+end
+
+
+function train_operators!(pl::PlannerLoss, N_epoch::Int, opt)
+        errors = []
+        for i = 1:N_epoch
+                data = data_gen()
+                for op in pl.operators
+                        train_operators!(op,data,opt)
+                end
+        end
+        return errors
+end
+
+
+#Construct and manage data
+abstract type DataConstructor end
+
+struct UniformDataConstructor <: DataConstructor
+        N::UInt64
+        satellites_bottom::Float32
+        satellites_top::Float32
+        debris_bottom::Float32
+        debris_top::Float32
+end
+function (dc::UniformDataConstructor)()
+        #currently ignores the quantity of data it should construct.
+    State(
+        rand(dc.satellites_bottom:dc.satellites_top, N_constellations)
+        , rand(dc.debris_bottom:dc.debris_top, N_debris)
+    )
+end
+
+end # end module

julia_code/Module/src/Orbits.jl

@@ -1,19 +1,22 @@
 module Orbits
 
-using Flux,LinearAlgebra,Zygote
+using Flux,LinearAlgebra,Zygote,BSON
 
 include("physical_model.jl")
-using .PhysicsModule
+export State, state_to_tuple, state_transition, BasicPhysics, PhysicalModel
 
 include("benefit_functions.jl")
-using .BenefitFunctions
+export LinearModel, EconomicModel
 
 include("flux_helpers.jl")
-using .NNTools
-
-# Exports
-export GeneralizedLoss , OperatorLoss, PlannerLoss, UniformDataConstructor
+export operator_policy_function_generator,
+        value_function_generator,
+        BranchGenerator,
+        cross_linked_planner_policy_function_generator
         
+# Exports from below
+export GeneralizedLoss ,OperatorLoss ,PlannerLoss ,UniformDataConstructor,
+        train_planner!, train_operators!
 
 # Code
 
@@ -160,4 +163,5 @@ function (dc::UniformDataConstructor)()
         , rand(dc.debris_bottom:dc.debris_top, N_debris)
     )
 end
+
 end # end module

julia_code/Module/src/benefit_functions.jl

@@ -1,11 +1,5 @@
+#don't write as a module
 
-module BenefitFunctions
-
-export LinearModel, BenefitFunction, EconomicModel
-
-
-include("physical_model.jl")
-using .PhysicsModule: State
 
 #=
 Benefit Functions:
@@ -66,4 +60,3 @@ function (em::CRRA)(
         return (core-1)/(1-em.σ)
 end
 
-end #end module

julia_code/Module/src/flux_helpers.jl

@@ -1,12 +1,3 @@
-module NNTools
-
-export BranchGenerator, value_function_generator
-
-using Flux, BSON, Zygote
-
-include("physical_model.jl")
-using .PhysicsModule: State,tuple
-
 
 
 #= TupleDuplicator
@@ -16,25 +7,26 @@ struct TupleDuplicator
         n::Int
 end
 #make tuples consisting of copies of whatever was provided
+function (f::TupleDuplicator)(s::State)
+	st = state_to_tuple(s)
+    return f(st)
+	#BROKEN: Fails in test, but works in REPL.
+end
+
 (f::TupleDuplicator)(x) = tuple([deepcopy(x) for i=1:f.n]...)
-#do the same but for the case of states, coerce them to tuples first
-(f::TupleDuplicator)(x::State) = tuple([deepcopy(tuple(x)) for i=1:f.n]...)
 
 
 
-#= BranchGenerator
-This generates a policy function full of branches with the properly scaled sides
-=#
 struct BranchGenerator
-        n::UInt8 
-        #limit to 2^8 operators
+        n::UInt
 end
-function (b::BranchGenerator)(branch::Flux.Chain,join_fn::Function) 
+function (b::BranchGenerator)(branch::Flux.Parallel,join_fn::Function) 
         # used to deepcopy the branches and duplicate the inputs in the returned chain
         f = TupleDuplicator(b.n)
         
-        return Flux.Chain(f,Flux.Parallel(join_fn,f(branch)))
-        #note that it both uses and runs the tuple duplicator f
+        return Flux.Chain(state_to_tuple,f,Flux.Parallel(join_fn,f(branch)...))
+        #note that it destructures the state to a tuple, duplicates,
+		# and then passes to the parallelized functions.
 end
 
 
@@ -82,7 +74,11 @@ function cross_linked_planner_policy_function_generator(number_params=32)
 
 end
 
-function operator_policy_function_generator(number_params=32)
+function operator_policy_function_generator(
+	N_constellations::Int
+	,N_debris
+	,number_params=32
+	)
 	return Flux.Chain(
 	    Flux.Parallel(vcat
 		        #parallel joins together stocks and debris
@@ -103,4 +99,3 @@ function operator_policy_function_generator(number_params=32)
 
 end
 
-end

julia_code/Module/src/physical_model.jl

@@ -1,20 +1,14 @@
-module PhysicsModule
 
-#Add exports here
-export State ,tuple ,state_transition ,PhysicalModel ,BasicModel
 
 
-using Flux, LinearAlgebra
-using Zygote #included for saving and loading models
-
 #==#
 struct State
         stocks::Array{Float32}
         debris::Array{Float32}
 end
 
-function tuple(st::State)
-        return (st.stocks, st.debris)
+function state_to_tuple(s::State)
+        return (s.stocks ,s.debris)
 end
 
 ### Physics
@@ -22,7 +16,7 @@ end
 
 abstract type PhysicalModel end
         
-struct BasicModel <: PhysicalModel
+struct BasicPhysics <: PhysicalModel
         #rate at which debris hits satellites
         debris_collision_rate::Real
         #rate at which satellites of different constellations collide
@@ -37,7 +31,7 @@ struct BasicModel <: PhysicalModel
         launch_debris_ratio::Real
 end
 function state_transition(
-        physics::BasicModel 
+        physics::BasicPhysics 
         ,state::State 
         ,launches::Vector{Float32}
         )
@@ -71,9 +65,8 @@ end
 
 function survival(
                 state::State
-                ,physical_model::BasicModel
+                ,physical_model::BasicPhysics
                 )	
         return exp.(-(physical_model.satellite_collision_rates .+ physical_model.decay_rate) * state.stocks .- (physical_model.debris_collision_rate * state.debris))
 end
 
-end #end module

julia_code/Module/tests/test_NNTools.jl

@@ -0,0 +1,54 @@
+using Test,Flux
+
+include("../src/Orbits.jl")
+using .Orbits
+
+
+
+#=
+Structure:
+
+This is broken into three parts
+ - The test_interfaces module contains various tools used in testing
+ - The test_routines includes functions that setup, run, and teardown tests
+ - Everything else is a set of tests using the standard testing tools in Julia
+=#
+@testset "Overall" verbose=true begin
+@testset "TupleDuplicator" verbose=true begin
+    #Check if tuple duplicator duplicates something
+    td2 = Orbits.TupleDuplicator(2)
+    @test typeof(td2) <: Orbits.TupleDuplicator
+
+    @test td2(([1,2],[3])) == (([1,2],[3]),([1,2],[3]))
+
+    st = State([1.0,2],[3])
+    @test td2((st.stocks,st.debris)) == (([1.0,2],[3]),([1.0,2],[3]))
+
+    @test td2(state_to_tuple(st)) == (([1.0,2],[3]),([1.0,2],[3]))
+    
+    @test td2(st) == ((st.stocks,st.debris),(st.stocks,st.debris))
+end
+
+@testset "BranchGenerator" verbose=true begin
+    st = State([1.0,2],[3])
+    tp = state_to_tuple(st)
+    bg = BranchGenerator(2)
+
+    branch = Flux.Parallel(vcat, Dense(2,1),Dense(1,1))
+
+    branched = bg(branch,vcat)
+
+    #type tests
+    @test typeof(branch(tp)) <: Array{Float32}
+
+    @test_broken typeof(branched[2](tp)) <: Array{Float32}
+    #ISSUE: what am I really looking for here?
+
+    #Check behaviors of the 
+    @test branched[1](st) == state_to_tuple(st) #Evaluation is of the wrong approach
+    
+    
+
+end #branch generator
+
+end #overall testset

julia_code/Module/workspace.code-workspace

@@ -0,0 +1,8 @@
+{
+	"folders": [
+		{
+			"path": ".."
+		}
+	],
+	"settings": {}
+}