You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
109 lines
4.6 KiB
Python
109 lines
4.6 KiB
Python
import torch
|
|
import combined as c
|
|
|
|
"""
|
|
This module holds the neural networks I am going to use to estimate
|
|
the functions of interest.
|
|
"""
|
|
|
|
class LaunchFnEstimand(torch.nn.Module):
|
|
"""
|
|
This is used to estimate the launch function
|
|
"""
|
|
def __init__(self, state_tensor_size,layers_size,number_constellations):
|
|
super().__init__()
|
|
self.number_constellations = number_constellations
|
|
self.layers_size = layers_size
|
|
self.state_tensor_size = state_tensor_size
|
|
|
|
|
|
#Layers
|
|
self.linear_1 = torch.nn.Linear(in_features=state_tensor_size, out_features=layers_size)
|
|
self.relu = torch.nn.ReLU()
|
|
self.linear_3 = torch.nn.Linear(in_features=layers_size, out_features=layers_size)
|
|
self.linear_5 = torch.nn.Linear(in_features=layers_size, out_features=number_constellations)
|
|
|
|
|
|
def forward(self, input_values):
|
|
|
|
intermediate_values = self.relu(input_values) #states should be positive anyway.
|
|
intermediate_values = self.linear_1(intermediate_values)
|
|
intermediate_values = self.linear_3(intermediate_values)
|
|
intermediate_values = self.linear_5(intermediate_values)
|
|
intermediate_values = self.relu(intermediate_values) #launches are always positive
|
|
|
|
return intermediate_values
|
|
|
|
class PartialDerivativesEstimand(torch.nn.Module):
|
|
"""
|
|
This is used to estimate the partial derivatives of the value functions
|
|
"""
|
|
def __init__(self,batch_size, number_constellations, number_states, layer_size=12):
|
|
super().__init__()
|
|
self.batch_size = batch_size
|
|
self.number_constellations = number_constellations
|
|
self.number_states = number_states
|
|
self.layer_size = layer_size
|
|
|
|
|
|
#preprocess (single linear layer in case there is anything that needs to happen to all states)
|
|
self.preprocess = torch.nn.Sequential(
|
|
torch.nn.ReLU() #cleanup as states must be positive
|
|
,torch.nn.Linear(in_features = self.number_states, out_features=self.number_states)
|
|
)
|
|
#upscale to get the basic dimensionality correct. From (batch,State) to (batch, constellation, state). Includes a reshape
|
|
self.upsample = lambda x: torch.nn.Upsample(scale_factor=self.number_constellations)(x).view(self.batch_size
|
|
,self.number_constellations
|
|
,self.number_states)
|
|
|
|
#sequential steps
|
|
self.sequential = torch.nn.Sequential(
|
|
torch.nn.Linear(in_features=number_states, out_features=layer_size)
|
|
#who knows if a convolution might help here.
|
|
,torch.nn.Linear(in_features=layer_size, out_features=layer_size)
|
|
,torch.nn.Linear(in_features=layer_size, out_features=layer_size)
|
|
)
|
|
|
|
#reduce the feature axis to match expected results
|
|
self.feature_reduction = torch.nn.Linear(in_features=layer_size, out_features=number_states)
|
|
|
|
def forward(self, input_values):
|
|
#Note that the input values are just going to be the state variables
|
|
#TODO:check that input values match the prepared dimension?
|
|
|
|
#preprocess
|
|
intermediate = self.preprocess(input_values)
|
|
|
|
#upscale the input values
|
|
intermediate = self.upsample(intermediate)
|
|
|
|
#intermediate processing
|
|
intermediate = self.sequential(intermediate)
|
|
|
|
#reduce feature axis to match the expected number of partials
|
|
intermediate = self.feature_reduction(intermediate)
|
|
|
|
|
|
return intermediate
|
|
|
|
class EstimandNN(torch.nn.Module):
|
|
"""
|
|
This neural network takes the current states as input values and returns both
|
|
the partial derivatives of the value function and the launch function.
|
|
"""
|
|
def __init__(self, state_tensor_size,layers_size,number_constellations):
|
|
super().__init__()
|
|
|
|
#So, this next section constructs different layers within the NN
|
|
#sinlge linear section
|
|
pass
|
|
#TODO:verify these are correct
|
|
self.partials_estimator = PartialDerivativesEstimand(state_tensor_size,layers_size,number_constellations) #TODO
|
|
self.launch_estimator = LaunchFnEstimand(state_tensor_size,layers_size,number_constellations)
|
|
|
|
def forward(self, input_values):
|
|
pass
|
|
partials = self.partials_estimator(input_values)
|
|
launch = self.launch_estimator(input_values)
|
|
|
|
return c.EstimandInterface(partials,launch) |