






from math import log

from scipy.constants import Boltzmann, physical_constants

a_							=    300  # TEMPERATURE (unit = K)


def        _a  (     UpperCamelCase_					:				float    ,   UpperCamelCase_					:				float    ,   UpperCamelCase_					:				float    ,   ) ->   float:
						"""simple docstring"""







						if donor_conc <= 0:
												raise ValueError("Donor concentration should be positive"   )
						elif acceptor_conc <= 0:
												raise ValueError("Acceptor concentration should be positive"   )
						elif intrinsic_conc <= 0:
												raise ValueError("Intrinsic concentration should be positive"   )
						elif donor_conc <= intrinsic_conc:
												raise ValueError(
												    "Donor concentration should be greater than intrinsic concentration"   )
						elif acceptor_conc <= intrinsic_conc:
												raise ValueError(
												    "Acceptor concentration should be greater than intrinsic concentration"   )
						else:
												return (
												    Boltzmann
												    * T
												    * log((donor_conc * acceptor_conc) / intrinsic_conc**2   )
												    / physical_constants["electron volt"][0]
												)


if __name__ == "__main__":
	import doctest

	doctest.testmod()














import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD


class      lowercase__     (			nn.Module	):
	a_         =42
	a_         =42
	a_         =0.0
	a_         =1
	a_         =1
	a_         =True
	a_         =False
	a_         =False
	a_         =False
	a_         =jnp.floataa

	def 		UpperCAmelCase  (							self       )->					Optional[Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						[]
							lowerCAmelCase__       =						[]

							for i in range(self.num_layers       ):
													lowerCAmelCase__       =						self.in_channels if i == 0 else self.out_channels

													lowerCAmelCase__       =						FlaxResnetBlockaD(
													    in_channels=__UpperCAmelCase		,    out_channels=self.out_channels		,    dropout_prob=self.dropout		,    dtype=self.dtype		,    )
													resnets.append(__UpperCAmelCase       )

													lowerCAmelCase__       =						FlaxTransformeraDModel(
													    in_channels=self.out_channels		,    n_heads=self.num_attention_heads		,    d_head=self.out_channels // self.num_attention_heads		,    depth=1		,    use_linear_projection=self.use_linear_projection		,    only_cross_attention=self.only_cross_attention		,    use_memory_efficient_attention=self.use_memory_efficient_attention		,    dtype=self.dtype		,    )
													attentions.append(__UpperCAmelCase       )

							lowerCAmelCase__       =						resnets
							lowerCAmelCase__       =						attentions

							if self.add_downsample:
													lowerCAmelCase__       =						FlaxDownsampleaD(self.out_channels		,    dtype=self.dtype       )


	def __call__(							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase=True       )->					str:







							'''simple docstring'''





							lowerCAmelCase__       =						()

							for resnet, attn in zip(self.resnets		,    self.attentions       ):
													lowerCAmelCase__       =						resnet(__UpperCAmelCase		,    __UpperCAmelCase		,    deterministic=__UpperCAmelCase       )
													lowerCAmelCase__       =						attn(__UpperCAmelCase		,    __UpperCAmelCase		,    deterministic=__UpperCAmelCase       )
													output_states += (hidden_states,)

							if self.add_downsample:
													lowerCAmelCase__       =						self.downsamplers_a(__UpperCAmelCase       )
													output_states += (hidden_states,)

							return hidden_states, output_states





class      lowercase__     (			nn.Module	):
	a_         =42
	a_         =42
	a_         =0.0
	a_         =1
	a_         =True
	a_         =jnp.floataa

	def 		UpperCAmelCase  (							self       )->					Dict:







							'''simple docstring'''





							lowerCAmelCase__       =						[]

							for i in range(self.num_layers       ):
													lowerCAmelCase__       =						self.in_channels if i == 0 else self.out_channels

													lowerCAmelCase__       =						FlaxResnetBlockaD(
													    in_channels=__UpperCAmelCase		,    out_channels=self.out_channels		,    dropout_prob=self.dropout		,    dtype=self.dtype		,    )
													resnets.append(__UpperCAmelCase       )
							lowerCAmelCase__       =						resnets

							if self.add_downsample:
													lowerCAmelCase__       =						FlaxDownsampleaD(self.out_channels		,    dtype=self.dtype       )


	def __call__(							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase=True       )->					int:







							'''simple docstring'''





							lowerCAmelCase__       =						()

							for resnet in self.resnets:
													lowerCAmelCase__       =						resnet(__UpperCAmelCase		,    __UpperCAmelCase		,    deterministic=__UpperCAmelCase       )
													output_states += (hidden_states,)

							if self.add_downsample:
													lowerCAmelCase__       =						self.downsamplers_a(__UpperCAmelCase       )
													output_states += (hidden_states,)

							return hidden_states, output_states





class      lowercase__     (			nn.Module	):
	a_         =42
	a_         =42
	a_         =42
	a_         =0.0
	a_         =1
	a_         =1
	a_         =True
	a_         =False
	a_         =False
	a_         =False
	a_         =jnp.floataa

	def 		UpperCAmelCase  (							self       )->					int:







							'''simple docstring'''





							lowerCAmelCase__       =						[]
							lowerCAmelCase__       =						[]

							for i in range(self.num_layers       ):
													lowerCAmelCase__       =						self.in_channels if (i == self.num_layers - 1) else self.out_channels
													lowerCAmelCase__       =						self.prev_output_channel if i == 0 else self.out_channels

													lowerCAmelCase__       =						FlaxResnetBlockaD(
													    in_channels=resnet_in_channels + res_skip_channels		,    out_channels=self.out_channels		,    dropout_prob=self.dropout		,    dtype=self.dtype		,    )
													resnets.append(__UpperCAmelCase       )

													lowerCAmelCase__       =						FlaxTransformeraDModel(
													    in_channels=self.out_channels		,    n_heads=self.num_attention_heads		,    d_head=self.out_channels // self.num_attention_heads		,    depth=1		,    use_linear_projection=self.use_linear_projection		,    only_cross_attention=self.only_cross_attention		,    use_memory_efficient_attention=self.use_memory_efficient_attention		,    dtype=self.dtype		,    )
													attentions.append(__UpperCAmelCase       )

							lowerCAmelCase__       =						resnets
							lowerCAmelCase__       =						attentions

							if self.add_upsample:
													lowerCAmelCase__       =						FlaxUpsampleaD(self.out_channels		,    dtype=self.dtype       )


	def __call__(							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase=True       )->					Dict:







							'''simple docstring'''





							for resnet, attn in zip(self.resnets		,    self.attentions       ):
													# pop res hidden states
													lowerCAmelCase__       =						res_hidden_states_tuple[-1]
													lowerCAmelCase__       =						res_hidden_states_tuple[:-1]
													lowerCAmelCase__       =						jnp.concatenate((hidden_states, res_hidden_states)		,    axis=-1       )

													lowerCAmelCase__       =						resnet(__UpperCAmelCase		,    __UpperCAmelCase		,    deterministic=__UpperCAmelCase       )
													lowerCAmelCase__       =						attn(__UpperCAmelCase		,    __UpperCAmelCase		,    deterministic=__UpperCAmelCase       )

							if self.add_upsample:
													lowerCAmelCase__       =						self.upsamplers_a(__UpperCAmelCase       )

							return hidden_states





class      lowercase__     (			nn.Module	):
	a_         =42
	a_         =42
	a_         =42
	a_         =0.0
	a_         =1
	a_         =True
	a_         =jnp.floataa

	def 		UpperCAmelCase  (							self       )->					int:







							'''simple docstring'''





							lowerCAmelCase__       =						[]

							for i in range(self.num_layers       ):
													lowerCAmelCase__       =						self.in_channels if (i == self.num_layers - 1) else self.out_channels
													lowerCAmelCase__       =						self.prev_output_channel if i == 0 else self.out_channels

													lowerCAmelCase__       =						FlaxResnetBlockaD(
													    in_channels=resnet_in_channels + res_skip_channels		,    out_channels=self.out_channels		,    dropout_prob=self.dropout		,    dtype=self.dtype		,    )
													resnets.append(__UpperCAmelCase       )

							lowerCAmelCase__       =						resnets

							if self.add_upsample:
													lowerCAmelCase__       =						FlaxUpsampleaD(self.out_channels		,    dtype=self.dtype       )


	def __call__(							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase=True       )->					int:







							'''simple docstring'''





							for resnet in self.resnets:
													# pop res hidden states
													lowerCAmelCase__       =						res_hidden_states_tuple[-1]
													lowerCAmelCase__       =						res_hidden_states_tuple[:-1]
													lowerCAmelCase__       =						jnp.concatenate((hidden_states, res_hidden_states)		,    axis=-1       )

													lowerCAmelCase__       =						resnet(__UpperCAmelCase		,    __UpperCAmelCase		,    deterministic=__UpperCAmelCase       )

							if self.add_upsample:
													lowerCAmelCase__       =						self.upsamplers_a(__UpperCAmelCase       )

							return hidden_states





class      lowercase__     (			nn.Module	):
	a_         =42
	a_         =0.0
	a_         =1
	a_         =1
	a_         =False
	a_         =False
	a_         =jnp.floataa

	def 		UpperCAmelCase  (							self       )->					Union[str, Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						[
							    FlaxResnetBlockaD(
							        in_channels=self.in_channels		,    out_channels=self.in_channels		,    dropout_prob=self.dropout		,    dtype=self.dtype		,    )
							]

							lowerCAmelCase__       =						[]

							for _ in range(self.num_layers       ):
													lowerCAmelCase__       =						FlaxTransformeraDModel(
													    in_channels=self.in_channels		,    n_heads=self.num_attention_heads		,    d_head=self.in_channels // self.num_attention_heads		,    depth=1		,    use_linear_projection=self.use_linear_projection		,    use_memory_efficient_attention=self.use_memory_efficient_attention		,    dtype=self.dtype		,    )
													attentions.append(__UpperCAmelCase       )

													lowerCAmelCase__       =						FlaxResnetBlockaD(
													    in_channels=self.in_channels		,    out_channels=self.in_channels		,    dropout_prob=self.dropout		,    dtype=self.dtype		,    )
													resnets.append(__UpperCAmelCase       )

							lowerCAmelCase__       =						resnets
							lowerCAmelCase__       =						attentions


	def __call__(							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase=True       )->					Union[str, Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						self.resnets[0](__UpperCAmelCase		,    __UpperCAmelCase       )
							for attn, resnet in zip(self.attentions		,    self.resnets[1:]       ):
													lowerCAmelCase__       =						attn(__UpperCAmelCase		,    __UpperCAmelCase		,    deterministic=__UpperCAmelCase       )
													lowerCAmelCase__       =						resnet(__UpperCAmelCase		,    __UpperCAmelCase		,    deterministic=__UpperCAmelCase       )

							return hidden_states






