Coverage for encodermap/loss_functions/loss_functions.py: 69%
267 statements
« prev ^ index » next coverage.py v7.1.0, created at 2023-02-07 11:05 +0000
1# -*- coding: utf-8 -*-
2# encodermap/loss_functions/loss_functions.py
3################################################################################
4# Encodermap: A python library for dimensionality reduction.
5#
6# Copyright 2019-2022 University of Konstanz and the Authors
7#
8# Authors:
9# Kevin Sawade, Tobias Lemke
10#
11# Encodermap is free software: you can redistribute it and/or modify
12# it under the terms of the GNU Lesser General Public License as
13# published by the Free Software Foundation, either version 2.1
14# of the License, or (at your option) any later version.
15# This package is distributed in the hope that it will be useful to other
16# researches. IT DOES NOT COME WITH ANY WARRANTY WHATSOEVER; without even the
17# implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
18# See the GNU Lesser General Public License for more details.
19#
20# See <http://www.gnu.org/licenses/>.
21################################################################################
22"""Loss functions for encodermap
24ToDo:
25 * Debug Autograph for distance cost
26 * WARNING: AutoGraph could not transform <function sigmoid_loss at 0x00000264AB761040> and will run it as-is.
27 * Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.
28 * Cause: module 'gast' has no attribute 'Index'
29 * To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert
31"""
32##############################################################################
33# Imports
34##############################################################################
36import tensorflow as tf
37import tensorflow.keras.backend as K
39from ..encodermap_tf1.misc import distance_cost
40from ..misc.distances import (
41 pairwise_dist,
42 pairwise_dist_periodic,
43 periodic_distance,
44 sigmoid,
45)
46from ..parameters.parameters import ADCParameters, Parameters, ParametersFramework
48##############################################################################
49# Globals
50##############################################################################
52__all__ = [
53 "reconstruction_loss",
54 "auto_loss",
55 "center_loss",
56 "regularization_loss",
57 "loss_combinator",
58 "distance_loss",
59 "cartesian_loss",
60 "cartesian_distance_loss",
61 "angle_loss",
62 "dihedral_loss",
63]
65##############################################################################
66# Functions for tf.cond
67# Don't know if this is really faster than logging every step to tensorboard
68##############################################################################
71def _do_nothing(*args):
72 """This function does nothing. One of the functions provided to tf.cond."""
73 pass
76def _summary_cost(name, cost):
77 """This functions logs a scalar to a name. One of the functions provided to tf.cond."""
78 tf.summary.scalar(name, cost)
81##############################################################################
82# Legacy Code to make some tests
83##############################################################################
86def old_distance_loss(model, parameters=None):
87 # choose parameters
88 if parameters is None: 88 ↛ 91line 88 didn't jump to line 91, because the condition on line 88 was never false
89 p = ParametersFramework(Parameters.defaults)
90 else:
91 p = parameters
92 # check Layers
93 if len(model.layers) == 2: 93 ↛ 95line 93 didn't jump to line 95, because the condition on line 93 was never true
94 # sequential API
95 latent = model.encoder
96 else:
97 # functional API
98 latent = model.encoder
100 # closure
101 def loss(y_true, y_pred=None, step=None):
102 loss.name = "distance_loss"
103 y_pred = latent(y_true, training=True)
104 # print(f'For testing: Loss latent of model to test em.encodermap_tf1.misc: {y_pred}')
105 if p.distance_cost_scale is not None: 105 ↛ 112line 105 didn't jump to line 112, because the condition on line 105 was never false
106 dist_cost = distance_cost(
107 y_true, y_pred, *p.dist_sig_parameters, p.periodicity
108 )
109 if p.distance_cost_scale != 0: 109 ↛ 113line 109 didn't jump to line 113, because the condition on line 109 was never false
110 dist_cost *= p.distance_cost_scale
111 else:
112 dist_cost = 0
113 tf.summary.scalar("Distance Cost", dist_cost)
114 return dist_cost
116 return loss
119##############################################################################
120# Public Functions
121##############################################################################
124def basic_loss_combinator(*losses):
125 """Calculates the sum of a list of losses and returns a combined loss.
127 The basic loss combinator does not write to summary. Can be used for debugging.
129 """
131 def loss(y_true, y_pred=None):
132 return sum([loss(y_true, y_pred) for loss in losses])
134 return loss
137def loss_combinator(*losses):
138 """Calculates the sum of a list of losses and returns a combined loss.
140 Args:
141 *losses: Variable length argument list of loss functions.
143 Returns:
144 function: A combined loss function that can be used in custom training or with model.fit()
146 Example:
147 >>> import encodermap as em
148 >>> from encodermap import loss_functions
149 >>> import tensorflow as tf
150 >>> import numpy as np
151 >>> tf.random.set_seed(1) # fix random state to pass doctest :)
153 >>> model = tf.keras.Sequential([
154 ... tf.keras.layers.Dense(100, kernel_regularizer=tf.keras.regularizers.l2(), activation='relu'),
155 ... tf.keras.layers.Dense(2, kernel_regularizer=tf.keras.regularizers.l2(), activation='relu'),
156 ... tf.keras.layers.Dense(100, kernel_regularizer=tf.keras.regularizers.l2(), activation='relu')
157 ... ])
159 >>> # Set up losses and bundle them using the loss combinator
160 >>> auto_loss = loss_functions.auto_loss(model)
161 >>> reg_loss = loss_functions.regularization_loss(model)
162 >>> loss = loss_functions.loss_combinator(auto_loss, reg_loss)
164 >>> # Compile model, model.fit() usually takes a tuple of (data, classes) but in
165 >>> # regression learning the data needs to be provided twice. That's why we use fit(data, data)
166 >>> model.compile(tf.keras.optimizers.Adam(), loss=loss)
167 >>> data = np.random.random((100, 100))
168 >>> history = model.fit(data, data, verbose=0)
169 >>> tf.random.set_seed(None) # reset seed
171 >>> # This weird contraption is also there to make the output predictable and pass tests
172 >>> # Somehow the tf.random.seed(1) does not work here. :(
173 >>> loss = history.history['loss'][0]
174 >>> print(loss) # doctest: +SKIP
175 {'loss': array([2.6])}
176 >>> print(type(loss))
177 <class 'float'>
179 """
181 def combined_loss_func(y_true, y_pred=None):
182 cost = sum([loss(y_true, y_pred) for loss in losses])
183 tf.summary.scalar("Combined Cost", cost)
184 return cost
186 return combined_loss_func
189def distance_loss(model, parameters=None, callback=None):
190 """Encodermap distance_loss
192 Transforms space using sigmoid function first proposed by sketch-map.
194 Args:
195 model (tf.keras.Model): A model you want to use the loss function on.
196 parameters (Union[encodermap.Parameters, None], optional): The parameters. If None is
197 provided default values (check them with print(em.Parameters.defaults_description()))
198 are used. Defaults to None.
200 Note:
201 If the model contains two layers. The first layer will be assumed to be the decoder.
202 If the model contains more layers, one layer needs to be named 'latent' (case insensitive).
204 Raises:
205 Exception: When no bottleneck/latent layer can be found in the model.
207 Returns:
208 function: A loss function.
210 References::
212 @article{ceriotti2011simplifying,
213 title={Simplifying the representation of complex free-energy landscapes using sketch-map},
214 author={Ceriotti, Michele and Tribello, Gareth A and Parrinello, Michele},
215 journal={Proceedings of the National Academy of Sciences},
216 volume={108},
217 number={32},
218 pages={13023--13028},
219 year={2011},
220 publisher={National Acad Sciences}
221 }
223 """
224 # choose parameters
225 if parameters is None:
226 p = ParametersFramework(Parameters.defaults)
227 else:
228 p = parameters
230 # check Layers
231 if len(model.layers) == 2:
232 # sequential API
233 latent = model.encoder
234 else:
235 # functional API
236 latent = model.encoder
237 # functional API without multiple models to be removed
238 # layer_index = [layer.name.lower() for layer in model.layers].index('latent')
239 # latent = tf.keras.Model(inputs=model.inputs, outputs=model.layers[layer_index].output)
241 if callback is None: 241 ↛ 244line 241 didn't jump to line 244, because the condition on line 241 was never false
242 write_bool = K.constant(False, "bool", name="log_bool")
243 else:
244 write_bool = callback.log_bool
246 # define dist loss
247 dist_loss = sigmoid_loss(p)
249 # closure
250 def distance_loss_func(y_true, y_pred=None):
251 """y_true can be whatever input you like, dihedrals, angles, pairwise dist, contact maps. That will be
252 transformed with Sketchmap's sigmoid function, as will the output of the latent layer of the autoencoder.
253 the difference of these two will result in a loss function."""
254 distance_loss_func.name = "distance_loss"
255 y_pred = latent(y_true, training=True)
256 # functional model gives a tuple
257 if isinstance(y_true, tuple): 257 ↛ 258line 257 didn't jump to line 258, because the condition on line 257 was never true
258 y_true = tf.concat(y_true, axis=1)
259 if p.distance_cost_scale is not None: 259 ↛ 264line 259 didn't jump to line 264, because the condition on line 259 was never false
260 dist_cost = dist_loss(y_true, y_pred)
261 if p.distance_cost_scale != 0: 261 ↛ 265line 261 didn't jump to line 265, because the condition on line 261 was never false
262 dist_cost *= p.distance_cost_scale
263 else:
264 dist_cost = 0
265 tf.cond(
266 write_bool,
267 true_fn=lambda: _summary_cost("Distance Cost", dist_cost),
268 false_fn=lambda: _do_nothing(),
269 name="Cost",
270 )
271 return dist_cost
273 return distance_loss_func
276def sigmoid_loss(parameters=None, periodicity_overwrite=None):
277 """Sigmoid loss closure for use in distance cost and cartesian distance cost.
279 Outer function prepares callable sigmoid. Sigmoid can then be called with just y_true and y_pred.
281 Args:
282 parameters (Union[encodermap.Parameters, None], optional): The parameters. If None is
283 provided default values (check them with print(em.Parameters.defaults_description()))
284 are used. Defaults to None.
285 periodicity overwrite(Union[float, None]), optional): Cartesian distance cost is always non-periodic.
286 To make sure no periodicity is applied to the data, set periodicity_overwrite to float('inf'). If
287 None is provided the periodicity of the parameters class (default 2*pi) will be used.
288 Defaults to None.
290 Returns:
291 function: A function that takes y_true and y_pred. Both need to be of the same shape.
293 """
294 if parameters is None:
295 p = ParametersFramework(Parameters.defaults)
296 else:
297 p = parameters
299 if periodicity_overwrite is not None:
300 periodicity = periodicity_overwrite
301 else:
302 periodicity = p.periodicity
304 # @tf.autograph.experimental.do_not_convert
305 def sigmoid_loss_func(y_true, y_pred):
306 r_h = y_true
307 r_l = y_pred
308 if periodicity == float("inf"):
309 dist_h = pairwise_dist(r_h)
310 else:
311 dist_h = pairwise_dist_periodic(r_h, periodicity)
312 dist_l = pairwise_dist(r_l)
314 sig_h = sigmoid(*p.dist_sig_parameters[:3])(dist_h)
315 sig_l = sigmoid(*p.dist_sig_parameters[3:])(dist_l)
317 cost = tf.reduce_mean(tf.square(sig_h - sig_l))
318 return cost
320 return sigmoid_loss_func
323def center_loss(model, parameters=None, callback=None):
324 """Encodermap center_loss
326 Use in custom training loops or in model.fit() training.
328 Args:
329 model (tf.keras.Model): A model you want to use the loss function on.
330 parameters (Union[encodermap.Parameters, None], optional): The parameters. If None is
331 provided default values (check them with print(em.Parameters.defaults_description()))
332 are used. Defaults to None.
334 Note:
335 If the model contains two layers. The first layer will be assumed to be the decoder.
336 If the model contains more layers, one layer needs to be named 'latent' (case insensitive).
338 Raises:
339 Exception: When no bottleneck/latent layer can be found in the model.
341 Returns:
342 function: A loss function.
344 """
345 # choose parameters
346 if parameters is None:
347 p = ParametersFramework(Parameters.defaults)
348 else:
349 p = parameters
350 # check Layers
351 if len(model.layers) == 2:
352 # sequential API
353 latent = model.encoder
354 else:
355 # functional API
356 latent = model.encoder
358 if callback is None: 358 ↛ 361line 358 didn't jump to line 361, because the condition on line 358 was never false
359 write_bool = K.constant(False, "bool", name="log_bool")
360 else:
361 write_bool = callback.log_bool
363 # closure
364 def center_loss_func(y_true, y_pred=None):
365 """y_true will not be used in this loss function. y_pred can be supplied, but if None will be taken from the
366 latent layer. This loss function tries to center the points in the latent layer."""
367 center_loss_func.name = "center_loss"
368 y_pred = latent(y_true, training=True)
369 # functional model gives a tuple
370 if isinstance(y_true, tuple): 370 ↛ 371line 370 didn't jump to line 371, because the condition on line 370 was never true
371 y_true = tf.concat(y_true, axis=1)
372 # center cost
373 # this is still a bit finicky
374 # needs to have tf.GradentTape() context manager to watch a single layer
375 if p.center_cost_scale is not None: 375 ↛ 385line 375 didn't jump to line 385, because the condition on line 375 was never false
376 # custom model dep code
377 # center_cost = tf.reduce_mean(tf.square(model.get_layer('Encoder')(y_true)))
378 # sequential model
379 # get the output according to https://keras.io/getting_started/faq/#how-can-i-obtain-the-output-of-an-intermediate-layer-feature-extraction
380 # this can be omitted by implementing a custom metric. Todo.
381 center_cost = tf.reduce_mean(tf.square(y_pred))
382 if p.center_cost_scale != 0: 382 ↛ 386line 382 didn't jump to line 386, because the condition on line 382 was never false
383 center_cost *= p.center_cost_scale
384 else:
385 center_cost = 0
386 tf.cond(
387 write_bool,
388 true_fn=lambda: _summary_cost("Center Cost", center_cost),
389 false_fn=lambda: _do_nothing(),
390 name="Cost",
391 )
392 return center_cost
394 return center_loss_func
397def regularization_loss(model, parameters=None, callback=None):
398 """Regularization loss of arbitrary tf.keras.Model
400 Use in custom training loops or in model.fit() training.
401 Loss is obtained as tf.math.add_n(model.losses)
403 Args:
404 model (tf.keras.Model): A model you want to use the loss function on.
406 Returns:
407 function: A loss function.
409 """
410 if parameters is None:
411 p = ParametersFramework(Parameters.defaults)
412 else:
413 p = parameters
415 if callback is None: 415 ↛ 418line 415 didn't jump to line 418, because the condition on line 415 was never false
416 write_bool = K.constant(False, "bool", name="log_bool")
417 else:
418 write_bool = callback.log_bool
420 def regularization_loss_func(y_true=None, y_pred=None):
421 """y_true and y_pred will not be considered here, because the regularization loss is accessed via model.losses."""
422 regularization_loss.name = "regularization_loss"
423 reg_loss = tf.math.add_n(model.losses)
424 tf.cond(
425 write_bool,
426 true_fn=lambda: _summary_cost("Regularization Cost", reg_loss),
427 false_fn=lambda: _do_nothing(),
428 name="Cost",
429 )
430 return reg_loss
432 return regularization_loss_func
435def reconstruction_loss(model):
436 """Simple Autoencoder recosntruction loss.
438 Use in custom training loops or in model.fit training.
440 Args:
441 model (tf.keras.Model): A model you want to use the loss function on.
443 Returns:
444 function: A loss function to be used in custom training or model.fit.
445 Function takes the following arguments:
446 y_true (tf.Tensor): The true tensor.
447 y_pred (tf.Tensor, optional): The output tensor. If not supplied
448 the model will be called to get this tensor. Defaults to None.
449 step (int): A step for tensorboard callbacks. Defaults to None.
451 Examples:
452 >>> import tensorflow as tf
453 >>> import encodermap as em
454 >>> from encodermap import loss_functions
455 >>> model = tf.keras.Model()
456 >>> loss = loss_functions.reconstruction_loss(model)
457 >>> x = tf.random.normal(shape=(10, 10))
458 >>> loss(x, x).numpy()
459 0.0
461 """
463 def reconstruction_loss_func(y_true, y_pred=None):
464 # if y_pred is None, this function is used in custom training
465 # and should use model to get the output
466 if y_pred is None: 466 ↛ 467line 466 didn't jump to line 467, because the condition on line 466 was never true
467 y_pred = model(y_true)
468 # calculate error
469 reconstruction_error = tf.reduce_mean(tf.square(tf.subtract(y_pred, y_true)))
470 return reconstruction_error
472 return reconstruction_loss_func
475def auto_loss(model, parameters=None, callback=None):
476 """Encodermap auto_loss.
478 Use in custom training loops or in model.fit() training.
480 Args:
481 model (tf.keras.Model): A model you want to use the loss function on.
482 parameters (Union[encodermap.Parameters, None], optional): The parameters. If None is
483 provided default values (check them with print(em.Parameters.defaults_description()))
484 are used. Defaults to None.
486 Returns:
487 function: A loss function.
489 """
490 if parameters is None:
491 p = ParametersFramework(Parameters.defaults)
492 else:
493 p = parameters
495 if callback is None: 495 ↛ 498line 495 didn't jump to line 498, because the condition on line 495 was never false
496 write_bool = K.constant(False, "bool", name="log_bool")
497 else:
498 write_bool = callback.log_bool
500 def auto_loss_func(y_true, y_pred=None):
501 """y_true is complete model input, y_pred is complete model output. Because here it is not intended to unpack
502 the output into dihedrals and angles, y_pred can be None and will be directly taken from the model."""
503 auto_loss_func.name = "auto_loss"
505 if y_pred is None:
506 y_pred = model(y_true)
507 if p.auto_cost_scale is not None: 507 ↛ 527line 507 didn't jump to line 527, because the condition on line 507 was never false
508 if p.auto_cost_variant == "mean_square": 508 ↛ 509line 508 didn't jump to line 509, because the condition on line 508 was never true
509 auto_cost = tf.reduce_mean(
510 tf.square(periodic_distance(y_true, y_pred, p.periodicity))
511 )
512 elif p.auto_cost_variant == "mean_abs": 512 ↛ 516line 512 didn't jump to line 516, because the condition on line 512 was never false
513 auto_cost = tf.reduce_mean(
514 tf.abs(periodic_distance(y_true, y_pred, p.periodicity))
515 )
516 elif p.auto_cost_variant == "mean_norm":
517 auto_cost = tf.reduce_mean(
518 tf.norm(periodic_distance(y_true, y_pred, p.periodicity), axis=1)
519 )
520 else:
521 raise ValueError(
522 "auto_cost_variant {} not available".format(p.auto_cost_variant)
523 )
524 if p.auto_cost_scale != 0: 524 ↛ 528line 524 didn't jump to line 528, because the condition on line 524 was never false
525 auto_cost *= p.auto_cost_scale
526 else:
527 auto_cost = 0
528 tf.cond(
529 write_bool,
530 true_fn=lambda: _summary_cost("Auto Cost", auto_cost),
531 false_fn=lambda: _do_nothing(),
532 name="Cost",
533 )
534 return auto_cost
536 return auto_loss_func
539def dihedral_loss(model, parameters=None, callback=None):
540 """Encodermap dihedral loss. Calculates distances between true and predicted dihedral angles.
542 Respects periodicity in a [-a, a] interval if the provided parameters have a periodicity of 2 * a.
544 Note:
545 The interval should be (-a, a], but due to floating point precision we can't make this
546 distinction here.
548 Args:
549 model (tf.keras.Model): The model to use the loss function on.
550 parameters (Union[encodermap.ADCParameters, None], optional): The parameters. If None is
551 provided default values (check them with print(em.ADCParameters.defaults_description()))
552 are used. Defaults to None.
554 Returns:
555 function: A loss function. Can be used in either custom training or model.fit().
557 """
558 if parameters is None: 558 ↛ 559line 558 didn't jump to line 559, because the condition on line 558 was never true
559 p = ParametersFramework(ADCParameters.defaults)
560 else:
561 p = parameters
563 if callback is None: 563 ↛ 566line 563 didn't jump to line 566, because the condition on line 563 was never false
564 write_bool = K.constant(False, "bool", name="log_bool")
565 else:
566 write_bool = callback.log_bool
568 # closure
569 def dihedral_loss_func(y_pred, y_true=None):
570 """y_pred should be model input dihedrals, y_true should be model output dihedrals."""
571 dihedral_loss_func.name = "dihedral_loss"
572 if p.dihedral_cost_scale is not None: 572 ↛ 593line 572 didn't jump to line 593, because the condition on line 572 was never false
573 if p.dihedral_cost_variant == "mean_square": 573 ↛ 574line 573 didn't jump to line 574, because the condition on line 573 was never true
574 dihedral_cost = tf.reduce_mean(
575 tf.square(periodic_distance(y_true, y_pred, p.periodicity))
576 )
577 elif p.dihedral_cost_variant == "mean_abs": 577 ↛ 581line 577 didn't jump to line 581, because the condition on line 577 was never false
578 dihedral_cost = tf.reduce_mean(
579 tf.abs(periodic_distance(y_true, y_pred, p.periodicity))
580 )
581 elif p.dihedral_cost_variant == "mean_norm":
582 dihedral_cost = tf.reduce_mean(
583 tf.norm(periodic_distance(y_true, y_pred, p.periodicity), axis=1)
584 )
585 else:
586 raise ValueError(
587 "dihedral_cost_variant {} not available".format(p.auto_cost_variant)
588 )
589 dihedral_cost /= p.dihedral_cost_reference
590 if p.dihedral_cost_scale != 0: 590 ↛ 594line 590 didn't jump to line 594, because the condition on line 590 was never false
591 dihedral_cost *= p.dihedral_cost_scale
592 else:
593 dihedral_cost = 0
594 tf.cond(
595 write_bool,
596 true_fn=lambda: _summary_cost("Dihedral Cost", dihedral_cost),
597 false_fn=lambda: _do_nothing(),
598 name="Cost",
599 )
600 return dihedral_cost
602 return dihedral_loss_func
605def side_dihedral_loss(model, parameters=None, callback=None):
606 """Encodermap sidechain dihedral loss. Calculates distances between true and predicted sidechain dihedral angles.
608 Respects periodicity in a [-a, a] interval if the provided parameters have a periodicity of 2 * a.
610 Note:
611 The interval should be (-a, a], but due to floating point precision we can't make this
612 distinction here.
614 Args:
615 model (tf.keras.Model): The model to use the loss function on.
616 parameters (Union[encodermap.ADCParameters, None], optional): The parameters. If None is
617 provided default values (check them with print(em.ADCParameters.defaults_description()))
618 are used. Defaults to None.
620 Returns:
621 function: A loss function. Can be used in either custom training or model.fit().
623 """
624 if parameters is None: 624 ↛ 625line 624 didn't jump to line 625, because the condition on line 624 was never true
625 p = ParametersFramework(ADCParameters.defaults)
626 else:
627 p = parameters
629 if callback is None: 629 ↛ 632line 629 didn't jump to line 632, because the condition on line 629 was never false
630 write_bool = K.constant(False, "bool", name="log_bool")
631 else:
632 write_bool = callback.log_bool
634 # closure
635 def side_dihedral_loss_func(y_pred, y_true=None):
636 """y_pred should be model input side dihedrals, y_true should be model output side dihedrals."""
637 side_dihedral_loss_func.name = "side_dihedral_loss"
638 if p.side_dihedral_cost_scale is not None:
639 if p.side_dihedral_cost_variant == "mean_square":
640 side_dihedral_cost = tf.reduce_mean(
641 tf.square(periodic_distance(y_true, y_pred, p.periodicity))
642 )
643 elif p.side_dihedral_cost_variant == "mean_abs":
644 side_dihedral_cost = tf.reduce_mean(
645 tf.abs(periodic_distance(y_true, y_pred, p.periodicity))
646 )
647 elif p.side_dihedral_cost_variant == "mean_norm":
648 side_dihedral_cost = tf.reduce_mean(
649 tf.norm(periodic_distance(y_true, y_pred, p.periodicity), axis=1)
650 )
651 else:
652 raise ValueError(
653 "dihedral_cost_variant {} not available".format(p.auto_cost_variant)
654 )
655 side_dihedral_cost /= p.side_dihedral_cost_reference
656 if p.side_dihedral_cost_scale != 0:
657 side_dihedral_cost *= p.side_dihedral_cost_scale
658 else:
659 side_dihedral_cost = 0
660 tf.cond(
661 write_bool,
662 true_fn=lambda: _summary_cost(
663 "Sidechain Dihedral Cost", side_dihedral_cost
664 ),
665 false_fn=lambda: _do_nothing(),
666 name="Cost",
667 )
668 return side_dihedral_cost
670 return side_dihedral_loss_func
673def angle_loss(model, parameters=None, callback=None):
674 """Encodermap angle loss. Calculates distances between true and predicted angles.
676 Respects periodicity in a [-a, a] interval if the provided parameters have a periodicity of 2 * a.
678 Note:
679 The interval should be (-a, a], but due to floating point precision we can't make this
680 distinction here.
682 Args:
683 model (tf.keras.Model): The model to use the loss function on.
684 parameters (Union[encodermap.ADCParameters, None], optional): The parameters. If None is
685 provided default values (check them with print(em.ADCParameters.defaults_description()))
686 are used. Defaults to None.
688 Returns:
689 function: A loss function. Can be used in either custom training or model.fit().
691 """
692 if parameters is None: 692 ↛ 693line 692 didn't jump to line 693, because the condition on line 692 was never true
693 p = ParametersFramework(ADCParameters.defaults)
694 else:
695 p = parameters
697 if callback is None: 697 ↛ 700line 697 didn't jump to line 700, because the condition on line 697 was never false
698 write_bool = K.constant(False, "bool", name="log_bool")
699 else:
700 write_bool = callback.log_bool
702 # closure
703 def angle_loss_func(y_pred, y_true=None):
704 """y_true should be input angles. y_pred should be output angles (either from mean input angles or, when
705 ADCParameters.use_backbone_angles == True, directly from model output)."""
706 angle_loss_func.name = "angle_loss"
707 if p.angle_cost_scale is not None: 707 ↛ 728line 707 didn't jump to line 728, because the condition on line 707 was never false
708 if p.angle_cost_variant == "mean_square": 708 ↛ 709line 708 didn't jump to line 709, because the condition on line 708 was never true
709 angle_cost = tf.reduce_mean(
710 tf.square(periodic_distance(y_true, y_pred, p.periodicity))
711 )
712 elif p.angle_cost_variant == "mean_abs": 712 ↛ 716line 712 didn't jump to line 716, because the condition on line 712 was never false
713 angle_cost = tf.reduce_mean(
714 tf.abs(periodic_distance(y_true, y_pred, p.periodicity))
715 )
716 elif p.angle_cost_variant == "mean_norm":
717 angle_cost = tf.reduce_mean(
718 tf.norm(periodic_distance(y_true, y_pred, p.periodicity), axis=1)
719 )
720 else:
721 raise ValueError(
722 "angle_cost_variant {} not available".format(p.auto_cost_variant)
723 )
724 angle_cost /= p.angle_cost_reference
725 if p.angle_cost_scale != 0: 725 ↛ 726line 725 didn't jump to line 726, because the condition on line 725 was never true
726 angle_cost *= p.angle_cost_scale
727 else:
728 angle_cost = 0
729 tf.cond(
730 write_bool,
731 true_fn=lambda: _summary_cost("Angle Cost", angle_cost),
732 false_fn=lambda: _do_nothing(),
733 name="Cost",
734 )
735 return angle_cost
737 return angle_loss_func
740def cartesian_distance_loss(model, parameters=None, callback=None):
741 """Encodermap cartesian distance loss. Calculates sigmoid-weighted distances between pairwise cartesians and latent.
743 Uses sketch-map's sigmoid function to transform the high-dimensional space of the input and the
744 low-dimensional space of latent.
746 Make sure to provide the pairwise cartesian distances. The output of the latent will be compared to the input.
748 Note:
749 If the model contains two layers. The first layer will be assumed to be the decoder.
750 If the model contains more layers, one layer needs to be named 'latent' (case insensitive).
752 Args:
753 model (tf.keras.Model): The model to use the loss function on.
754 parameters (Union[encodermap.ADCParameters, None], optional): The parameters. If None is
755 provided default values (check them with print(em.ADCParameters.defaults_description()))
756 are used. Defaults to None.
758 Returns:
759 function: A loss function. Can be used in either custom training or model.fit().
761 """
762 if parameters is None: 762 ↛ 763line 762 didn't jump to line 763, because the condition on line 762 was never true
763 p = ParametersFramework(ADCParameters.defaults)
764 else:
765 p = parameters
767 if callback is None: 767 ↛ 770line 767 didn't jump to line 770, because the condition on line 767 was never false
768 write_bool = K.constant(False, "bool", name="log_bool")
769 else:
770 write_bool = callback.log_bool
772 dist_loss = sigmoid_loss(p, periodicity_overwrite=float("inf"))
774 def cartesian_distance_loss_func(y_true, y_pred):
775 """y_true can be whatever input you like, dihedrals, angles, pairwise dist, contact maps. That will be
776 transformed with Sketchmap's sigmoid function, as will the output of the latent layer of the autoencoder.
777 the difference of these two will result in a loss function."""
778 cartesian_distance_loss_func.name = "cartesian_distance_loss"
779 if p.cartesian_distance_cost_scale is not None: 779 ↛ 784line 779 didn't jump to line 784, because the condition on line 779 was never false
780 dist_cost = dist_loss(y_true, y_pred)
781 if p.distance_cost_scale != 0: 781 ↛ 785line 781 didn't jump to line 785, because the condition on line 781 was never false
782 dist_cost *= p.cartesian_distance_cost_scale
783 else:
784 dist_cost = 0
785 tf.cond(
786 write_bool,
787 true_fn=lambda: _summary_cost("Cartesian Distance Cost", dist_cost),
788 false_fn=lambda: _do_nothing(),
789 name="Cost",
790 )
791 return dist_cost
793 return cartesian_distance_loss_func
796def cartesian_loss(
797 model,
798 scale_callback=None,
799 parameters=None,
800 log_callback=None,
801 print_current_scale=False,
802):
803 """Encodermap cartesian distance loss. Calculates sigmoid-weighted distances between pairwise cartesians and latent.
805 Uses sketch-map's sigmoid function to transform the high-dimensional space of the input and the
806 high-dimensional space of the output.
808 Adjustments to this cost_function via the soft_start parameter need to be made via a callback that re-compiles the
809 model during training. For this, the soft_start parameters of the outer function will be used.
810 It must be either 0 or 1, indexing the 1st or 2nd element of the cartesian_cost_scale_soft_start
811 tuple. The callback should also be provided when model.fit is executed.
813 Three cases are possible:
814 * Case 1: step < cartesian_cost_scale_soft_start[0]: cost_scale = 0
815 * Case 2: cartesian_cost_scale_soft_start[0] <= step <= cartesian_cost_scale_soft_start[1]:
816 cost_scale = p.cartesian_cost_scale / (cartesian_cost_scale_soft_start[1] - cartesian_cost_scale_soft_start[0]) * step
817 * Case 3: cartesian_cost_scale_soft_start[1] < step: cost_scale = p.cartesian_cost_scale
819 Make sure to provide the pairwise cartesian distances. This function will be adjusted as training increases via a
820 callback. See encodermap.callbacks.callbacks.IncreaseCartesianCost for more info.
822 Args:
823 model (tf.keras.Model): The model to use the loss function on.
824 parameters (Union[encodermap.ADCParameters, None], optional): The parameters. If None is
825 provided default values (check them with print(em.ADCParameters.defaults_description()))
826 are used. Defaults to None.
827 soft_start (Union[int, None], optional): How to scale the cartesian loss. The ADCParameters class contains a
828 two-tuple of integers. These integers can be used to scale this loss function. If soft_start is 0, the
829 first value of ADCParameters.cartesian_cost_scale_soft_start will be used, if it is 1, the second.
830 if it is None, or both values of ADCParameters.cartesian_cost_scale_soft_start are None, the cost will
831 not be scaled. Defaults to None.
832 print_current_scale (bool, optional): Whether to print the current scale. Is used in unittesting. Defaults to False.
834 Raises:
835 Exception: When no bottleneck/latent layer can be found in the model.
836 Exception: When soft_start is greater than 1 and can't index the two-tuple.
838 Returns:
839 function: A loss function. Can be used in either custom training or model.fit().
841 """
842 if parameters is None: 842 ↛ 843line 842 didn't jump to line 843, because the condition on line 842 was never true
843 p = ParametersFramework(ADCParameters.defaults)
844 else:
845 p = parameters
846 if scale_callback is not None:
847 current_scale_callback = scale_callback.current_cartesian_cost_scale
848 else:
849 current_scale_callback = K.constant(
850 p.cartesian_cost_scale, dtype="float32", name="current_cartesian_cost_scale"
851 )
853 if print_current_scale:
854 print(current_scale_callback)
856 if log_callback is None: 856 ↛ 859line 856 didn't jump to line 859, because the condition on line 856 was never false
857 write_bool = K.constant(False, "bool", name="log_bool")
858 else:
859 write_bool = log_callback.log_bool
861 def cartesian_loss_func(y_true, y_pred=None):
862 """y_true should be pairwise distances of input cartesians,
863 y_pred should be pairwise distances of back-mapped output cartesians."""
864 scale = current_scale_callback
865 if p.cartesian_cost_variant == "mean_square": 865 ↛ 866line 865 didn't jump to line 866, because the condition on line 865 was never true
866 cartesian_cost = tf.reduce_mean(tf.square(y_true - y_pred))
867 elif p.cartesian_cost_variant == "mean_abs": 867 ↛ 869line 867 didn't jump to line 869, because the condition on line 867 was never false
868 cartesian_cost = tf.reduce_mean(tf.abs(y_true - y_pred))
869 elif p.cartesian_cost_variant == "mean_norm":
870 cartesian_cost = tf.reduce_mean(tf.norm(y_true - y_pred, axis=1))
871 else:
872 raise ValueError(
873 "cartesian_cost_variant {} not available".format(
874 p.dihedral_to_cartesian_cost_variant
875 )
876 )
877 cartesian_cost /= p.cartesian_cost_reference
878 tf.cond(
879 write_bool,
880 true_fn=lambda: _summary_cost(
881 "Cartesian Cost before scaling", cartesian_cost
882 ),
883 false_fn=lambda: _do_nothing(),
884 name="Cost",
885 )
886 tf.cond(
887 write_bool,
888 true_fn=lambda: _summary_cost("Cartesian Cost current scaling", scale),
889 false_fn=lambda: _do_nothing(),
890 name="Cost",
891 )
892 cartesian_cost *= scale
893 tf.cond(
894 write_bool,
895 true_fn=lambda: _summary_cost(
896 "Cartesian Cost after scaling", cartesian_cost
897 ),
898 false_fn=lambda: _do_nothing(),
899 name="Cost",
900 )
901 return cartesian_cost
903 return cartesian_loss_func